xbarray 0.0.1a13__py3-none-any.whl

array_api_typing/typing_extra/_at.py

@@ -0,0 +1,87 @@
+from typing import Protocol, TypeVar, Optional, Any, Tuple, Union, Type, TypedDict, List
+from abc import abstractmethod
+from array_api_typing.typing_compat._api_typing import _NAMESPACE_ARRAY
+from array_api_typing.typing_compat._array_typing import SetIndex
+
+class AtResult(Protocol[_NAMESPACE_ARRAY]):
+    @abstractmethod
+    def __getitem__(self, idx: SetIndex, /) -> "AtResult[_NAMESPACE_ARRAY]":
+        """
+        Allow for the alternate syntax ``at(x)[start:stop:step]``.
+
+        It looks prettier than ``at(x, slice(start, stop, step))``
+        and feels more intuitive coming from the JAX documentation.
+        """
+        pass
+
+    @abstractmethod
+    def set(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def add(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def subtract(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def multiply(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def divide(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def power(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+
+    @abstractmethod
+    def min(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass
+    
+    @abstractmethod
+    def max(
+        self,
+        y: Union[_NAMESPACE_ARRAY, float, int, complex],
+        /,
+        copy: Optional[bool] = None,
+    ) -> _NAMESPACE_ARRAY:
+        pass

xbarray/__init__.py

@@ -0,0 +1 @@
+from .backends.base import * # Import Abstract Typings

xbarray/backends/_cls_base.py

@@ -0,0 +1,9 @@
+from typing import Type, Generic
+from .base import BArrayType, BDeviceType, BDtypeType, BRNGType
+
+class ComputeBackendImplCls(Generic[BArrayType, BDeviceType, BDtypeType, BRNGType], Type):
+    def __str__(self):
+        return self.simplified_name
+    
+    def __repr__(self):
+        return self.simplified_name

xbarray/backends/_implementations/_common/implementations.py

@@ -0,0 +1,87 @@
+from typing import Any, Union, Callable, Mapping, Sequence
+from array_api_typing.typing_compat import ArrayAPINamespace as CompatNamespace, ArrayAPIArray as CompatArray, ArrayAPIDType as CompatDType
+import array_api_compat
+import dataclasses
+
+__all__ = [
+    "get_abbreviate_array_function",
+    "get_map_fn_over_arrays_function",
+    "get_pad_dim_function",
+]
+
+def get_abbreviate_array_function(
+    backend : CompatNamespace[CompatArray, Any, Any],
+    default_integer_dtype : CompatDType,
+    func_dtype_is_real_floating : Callable[[CompatDType], bool],
+    func_dtype_is_real_integer : Callable[[CompatDType], bool],
+    func_dtype_is_boolean : Callable[[CompatDType], bool],
+):
+    def abbreviate_array(array : CompatArray, try_cast_scalar : bool = True) -> Union[float, int, CompatArray]:
+        """
+        Abbreivates an array to a single element if possible.
+        Or, if some dimensions are the same, abbreviates to a smaller array (but with the same number of dimensions).
+        """
+        abbr_array = array
+        idx = backend.zeros(1, dtype=default_integer_dtype, device=array_api_compat.device(abbr_array))
+        for dim_i in range(len(array.shape)):
+            first_elem = backend.take(abbr_array, idx, axis=dim_i)
+            if backend.all(abbr_array == first_elem):
+                abbr_array = first_elem
+            else:
+                continue
+        if try_cast_scalar:
+            if all(i == 1 for i in abbr_array.shape):
+                elem = abbr_array[tuple([0] * len(abbr_array.shape))]
+                if func_dtype_is_real_floating(elem.dtype):
+                    return float(elem)
+                elif func_dtype_is_real_integer(elem.dtype):
+                    return int(elem)
+                elif func_dtype_is_boolean(elem.dtype):
+                    return bool(elem)
+                else:
+                    raise ValueError(f"Abbreviated array element dtype must be a real floating or integer or boolean type, actual dtype: {elem.dtype}")
+        else:
+            return array
+    return abbreviate_array
+
+def get_map_fn_over_arrays_function(
+    is_backendarray : Callable[[Any], bool],
+):
+    def map_fn_over_arrays(data : Any, func : Callable[[CompatArray], CompatArray]) -> Any:
+        """
+        Map a function to the data.
+        """
+        if is_backendarray(data):
+            return func(data)
+        elif isinstance(data, Mapping):
+            ret = {k: map_fn_over_arrays(v, func) for k, v in data.items()}
+            try:
+                return type(data)(**ret)  # try to keep the same mapping type
+            except:
+                return ret
+        elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
+            ret = [map_fn_over_arrays(i, func) for i in data]
+            try:
+                return type(data)(ret)  # try to keep the same sequence type
+            except:
+                return ret
+        elif dataclasses.is_dataclass(data):
+            return type(data)(**map_fn_over_arrays(dataclasses.asdict(data), func))
+        else:
+            return data
+    return map_fn_over_arrays
+
+def get_pad_dim_function(
+    backend : CompatNamespace[CompatArray, Any, Any],
+):
+    def pad_dim(x : CompatArray, dim : int, target_size : int, value : Union[float, int] = 0) -> CompatArray:
+        size_at_target = x.shape[dim]
+        if size_at_target > target_size:
+            raise ValueError(f"Cannot pad dimension {dim} of size {size_at_target} to smaller target size {target_size}")
+        if size_at_target == target_size:
+            return x
+        target_shape = list(x.shape)
+        target_shape[dim] = target_size - size_at_target
+        pad_value = backend.full(target_shape, value, dtype=x.dtype, device=array_api_compat.device(x))
+        return backend.concat([x, pad_value], axis=dim)
+    return pad_dim

xbarray/backends/_implementations/jax/__init__.py

@@ -0,0 +1,33 @@
+import jax.numpy
+
+if hasattr(jax.numpy, "__array_api_version__"):
+    compat_module = jax.numpy
+    from jax.numpy import *
+    from jax.numpy import __array_api_version__, __array_namespace_info__
+else:
+    import jax.experimental.array_api as compat_module
+    from jax.experimental.array_api import *
+    from jax.experimental.array_api import __array_api_version__, __array_namespace_info__
+
+from array_api_compat.common._helpers import *
+
+simplified_name = "jax"
+
+# Import and bind all functions from array_api_extra before exposing them
+import array_api_extra
+from functools import partial
+for api_name in dir(array_api_extra):
+    if api_name.startswith('_'):
+        continue
+
+    if api_name in ['at', 'broadcast_shapes']:
+        globals()[api_name] = getattr(array_api_extra, api_name)
+    else:
+        globals()[api_name] = partial(
+            getattr(array_api_extra, api_name),
+            xp=compat_module
+        )
+
+from ._typing import *
+from ._extra import *
+__import__(__package__ + ".random")

xbarray/backends/_implementations/jax/_extra.py

@@ -0,0 +1,127 @@
+from typing import Any, Union, Optional, Callable
+import jax
+import jax.numpy as jnp
+import numpy as np
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+from xbarray.backends.base import ComputeBackend, SupportsDLPack
+
+__all__ = [
+    "default_integer_dtype",
+    "default_index_dtype",
+    "default_floating_dtype",
+    "default_boolean_dtype",
+    "serialize_device",
+    "deserialize_device",
+    "is_backendarray",
+    "from_numpy",
+    "from_other_backend",
+    "to_numpy",
+    "to_dlpack",
+    "dtype_is_real_integer",
+    "dtype_is_real_floating",
+    "dtype_is_boolean",
+    "abbreviate_array",
+    "map_fn_over_arrays",
+    "pad_dim",
+]
+
+default_integer_dtype = int
+default_index_dtype = int
+default_floating_dtype = float
+default_boolean_dtype = bool
+
+def serialize_device(device : Optional[Any]) -> Optional[str]:
+    if device is None:
+        return None
+    assert hasattr(device, "platform") and hasattr(device, "id")
+    return device.platform + ":" + str(device.id)
+
+def deserialize_device(device_str : Optional[str]) -> Optional[Any]:
+    if device_str is None:
+        return None
+    platform, id_str = device_str.split(":")
+    id = int(id_str)
+    # Find the device with matching platform and id
+    for device in jax.devices(platform):
+        if device.platform == platform and device.id == id:
+            return device
+    # Fallback with the same platform
+    for device in jax.devices(platform):
+        if device.platform == platform:
+            return device
+    raise ValueError(f"Device with platform '{platform}' and id '{id}' not found.")
+
+def is_backendarray(data : Any) -> bool:
+    return isinstance(data, jax.Array)
+
+def from_numpy(
+    data : np.ndarray,
+    /,
+    *,
+    dtype : Optional[DTYPE_TYPE] = None,
+    device : Optional[DEVICE_TYPE] = None
+) -> ARRAY_TYPE:
+    return jax.numpy.asarray(data, dtype=dtype, device=device)
+
+def from_other_backend(
+    other_backend: ComputeBackend,
+    data: Any,
+    /,
+) -> ARRAY_TYPE:
+    data_dlpack = other_backend.to_dlpack(data)
+    return jax.dlpack.from_dlpack(data_dlpack)
+    # except Exception as e:
+    #     # jax sometimes has tiling issues with dlpack converted data
+    #     np = other_backend.to_numpy(data)
+    #     return from_numpy(np)
+
+def to_numpy(
+    data : ARRAY_TYPE
+) -> np.ndarray:
+    if data.dtype == jax.dtypes.bfloat16:
+        data = data.astype(np.float32)
+    return np.asarray(data)
+
+def to_dlpack(
+    data: ARRAY_TYPE,
+    /,
+) -> SupportsDLPack:
+    return data
+
+def dtype_is_real_integer(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return np.issubdtype(dtype, np.integer)
+
+def dtype_is_real_floating(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype == jax.dtypes.bfloat16 or np.issubdtype(dtype, np.floating)
+
+def dtype_is_boolean(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype == np.bool_ or dtype == bool
+
+from .._common.implementations import *
+if hasattr(jax.numpy, "__array_api_version__"):
+    compat_module = jax.numpy
+else:
+    import jax.experimental.array_api as compat_module
+abbreviate_array = get_abbreviate_array_function(
+    backend=compat_module,
+    default_integer_dtype=default_integer_dtype,
+    func_dtype_is_real_floating=dtype_is_real_floating,
+    func_dtype_is_real_integer=dtype_is_real_integer,
+    func_dtype_is_boolean=dtype_is_boolean
+)
+def map_fn_over_arrays(
+    data : Any, func : Callable[[ARRAY_TYPE], ARRAY_TYPE]
+):
+    return jax.tree.map(
+        func,
+        data
+    )
+pad_dim = get_pad_dim_function(
+    backend=compat_module,
+)

xbarray/backends/_implementations/jax/_typing.py

@@ -0,0 +1,15 @@
+from typing import Union, Any, Optional
+import jax
+import numpy as np
+
+__all__ = [
+    'ARRAY_TYPE',
+    'DTYPE_TYPE',
+    'DEVICE_TYPE',
+    'RNG_TYPE',
+]
+
+ARRAY_TYPE = jax.Array
+DTYPE_TYPE = np.dtype
+DEVICE_TYPE = Union[jax.Device, jax.sharding.Sharding]
+RNG_TYPE = jax.Array

xbarray/backends/_implementations/jax/random.py

@@ -0,0 +1,115 @@
+from typing import Union, Optional, Tuple, Any
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+import numpy as np
+import jax
+
+__all__ = [
+    "random_number_generator",
+    "random_discrete_uniform",
+    "random_uniform",
+    "random_exponential",
+    "random_normal",
+    "random_geometric",
+    "random_permutation"
+]
+
+def random_number_generator(
+    seed : Optional[int] = None,
+    *,
+    device : Optional[DEVICE_TYPE] = None
+) -> RNG_TYPE:
+    rng_seed = np.random.randint(65535) if seed is None else seed
+    rng = jax.random.key(
+        seed=rng_seed
+    )
+    return rng
+
+def random_discrete_uniform(
+    shape : Union[int, Tuple[int, ...]], 
+    /,
+    from_num : int, 
+    to_num : int, 
+    *,
+    rng : RNG_TYPE, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    t = jax.random.randint(rng, shape, minval=int(from_num), maxval=int(to_num), dtype=dtype or int)
+    if device is not None:
+        t = jax.device_put(t, device)
+    return new_rng, t
+
+def random_uniform(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE, 
+    low : float = 0.0, high : float = 1.0,
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    data = jax.random.uniform(rng, shape, dtype=dtype or float, minval=low, maxval=high)
+    if device is not None:
+        data = jax.device_put(data, device)
+    return new_rng, data
+
+def random_exponential(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE,  
+    lambd : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    data = jax.random.exponential(rng, shape, dtype=dtype or float) / lambd
+    if device is not None:
+        data = jax.device_put(data, device)
+    return new_rng, data
+
+def random_normal(
+    shape: Union[int, Tuple[int, ...]],
+    /,
+    *, 
+    rng : RNG_TYPE,
+    mean : float = 0.0, std : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    data = jax.random.normal(rng, shape, dtype=dtype or float) * std + mean
+    if device is not None:
+        data = jax.device_put(data, device)
+    return new_rng, data
+
+def random_geometric(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    p: float, 
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None, 
+    device: Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    data = jax.random.geometric(rng, p=p, shape=shape, dtype=dtype or int)
+    if device is not None:
+        data = jax.device_put(data, device)
+    return new_rng, data
+
+def random_permutation(
+    n : int,
+    /,
+    *,
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None,
+    device: Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    new_rng, rng = jax.random.split(rng)
+    data = jax.random.permutation(rng, n, dtype=dtype or int)
+    if device is not None:
+        data = jax.device_put(data, device)
+    return new_rng, data

xbarray/backends/_implementations/numpy/__init__.py

@@ -0,0 +1,25 @@
+from array_api_compat.numpy import *
+from array_api_compat.numpy import __array_api_version__, __array_namespace_info__
+from array_api_compat.common._helpers import *
+
+simplified_name = "numpy"
+
+from array_api_compat import numpy as compat_module
+# Import and bind all functions from array_api_extra before exposing them
+import array_api_extra
+from functools import partial
+for api_name in dir(array_api_extra):
+    if api_name.startswith('_'):
+        continue
+
+    if api_name in ['at', 'broadcast_shapes']:
+        globals()[api_name] = getattr(array_api_extra, api_name)
+    else:
+        globals()[api_name] = partial(
+            getattr(array_api_extra, api_name),
+            xp=compat_module
+        )
+
+from ._typing import *
+from ._extra import *
+__import__(__package__ + ".random")

xbarray/backends/_implementations/numpy/_extra.py

@@ -0,0 +1,98 @@
+from typing import Any, Union, Optional
+import numpy as np
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+from xbarray.backends.base import ComputeBackend, SupportsDLPack
+
+__all__ = [
+    "default_integer_dtype",
+    "default_index_dtype",
+    "default_floating_dtype",
+    "default_boolean_dtype",
+    "serialize_device",
+    "deserialize_device",
+    "is_backendarray",
+    "from_numpy",
+    "from_other_backend",
+    "to_numpy",
+    "to_dlpack",
+    "dtype_is_real_integer",
+    "dtype_is_real_floating",
+    "dtype_is_boolean",
+    "abbreviate_array",
+    "map_fn_over_arrays",
+    "pad_dim",
+]
+
+default_integer_dtype = int
+default_index_dtype = int
+default_floating_dtype = float
+default_boolean_dtype = bool
+
+def serialize_device(device : Optional[DEVICE_TYPE]) -> Optional[str]:
+    return None  # NumPy does not have device concept
+
+def deserialize_device(device_str : Optional[str]) -> Optional[DEVICE_TYPE]:
+    return None  # NumPy does not have device concept
+
+def is_backendarray(data : Any) -> bool:
+    return isinstance(data, np.ndarray)
+
+def from_numpy(
+    data : np.ndarray,
+    /,
+    *,
+    dtype : Optional[DTYPE_TYPE] = None,
+    device : Optional[DEVICE_TYPE] = None
+) -> ARRAY_TYPE:
+    return data
+
+def from_other_backend(
+    other_backend: ComputeBackend,
+    data: Any,
+    /,
+) -> ARRAY_TYPE:
+    return other_backend.to_numpy(data)
+
+def to_numpy(
+    data : ARRAY_TYPE
+) -> np.ndarray:
+    return data
+
+def to_dlpack(
+    data: ARRAY_TYPE,
+    /,
+) -> SupportsDLPack:
+    return data
+
+def dtype_is_real_integer(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return np.issubdtype(dtype, np.integer)
+
+def dtype_is_real_floating(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return np.issubdtype(dtype, np.floating)
+
+def dtype_is_boolean(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype == np.bool_ or dtype == bool
+
+from .._common.implementations import *
+from array_api_compat import numpy as compat_module
+abbreviate_array = get_abbreviate_array_function(
+    backend=compat_module,
+    default_integer_dtype=default_integer_dtype,
+    func_dtype_is_real_floating=dtype_is_real_floating,
+    func_dtype_is_real_integer=dtype_is_real_integer,
+    func_dtype_is_boolean=dtype_is_boolean,
+)
+
+map_fn_over_arrays = get_map_fn_over_arrays_function(
+    is_backendarray=is_backendarray,
+)
+
+pad_dim = get_pad_dim_function(
+    backend=compat_module,
+)

xbarray/backends/_implementations/numpy/_typing.py

@@ -0,0 +1,14 @@
+from typing import Union, Any, Optional
+import numpy as np
+
+__all__ = [
+    'ARRAY_TYPE',
+    'DTYPE_TYPE',
+    'DEVICE_TYPE',
+    'RNG_TYPE',
+]
+
+ARRAY_TYPE = np.ndarray
+DTYPE_TYPE = np.dtype
+DEVICE_TYPE = Any
+RNG_TYPE = np.random.Generator