xbarray 0.0.1a1__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 .base import *

xbarray/base/__init__.py

@@ -0,0 +1 @@
+from .base import *

xbarray/base/base.py

@@ -0,0 +1,199 @@
+from typing import Optional, Generic, TypeVar, Dict, Union, Any, Sequence, SupportsFloat, Tuple, Type, Callable, Mapping, Protocol
+import abc
+from array_api_typing.typing_extra import *
+import numpy as np
+
+ArrayAPISetIndex = SetIndex
+ArrayAPIGetIndex = GetIndex
+
+__all__ = [
+    "RNGBackend",
+    "ComputeBackend",
+    "BArrayType",
+    "BDeviceType",
+    "BDtypeType",
+    "BRNGType",
+    "SupportsDLPack",
+    "ArrayAPIArray",
+    "ArrayAPIDevice",
+    "ArrayAPIDType",
+    "ArrayAPINamespace",
+    "ArrayAPISetIndex",
+    "ArrayAPIGetIndex",
+]
+
+BArrayType = TypeVar("BArrayType", covariant=True, bound=ArrayAPIArray)
+BDeviceType = TypeVar("BDeviceType", covariant=True, bound=ArrayAPIDevice)
+BDtypeType = TypeVar("BDtypeType", covariant=True, bound=ArrayAPIDType)
+BRNGType = TypeVar("BRNGType", covariant=True)
+class RNGBackend(Protocol[BArrayType, BDeviceType, BDtypeType, BRNGType]):
+    @abc.abstractmethod
+    def random_number_generator(
+        self, 
+        seed : Optional[int] = None, 
+        *,
+        device : Optional[BDeviceType] = None
+    ) -> BRNGType:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def random_discrete_uniform(
+        self, 
+        shape : Union[int, Tuple[int, ...]], 
+        from_num : int, 
+        to_num : int, 
+        /,
+        *,
+        rng : BRNGType, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None, 
+    ) -> Tuple[BRNGType, BArrayType]:
+        """
+        Sample from a discrete uniform distribution [from_num, to_num) with shape `shape`.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def random_uniform(
+        self, 
+        shape: Union[int, Tuple[int, ...]], 
+        /,
+        *,
+        rng : BRNGType, 
+        low : float = 0.0, high : float = 1.0, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None,
+    ) -> Tuple[BRNGType, BArrayType]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def random_exponential(
+        self, 
+        shape: Union[int, Tuple[int, ...]], 
+        /,
+        *,
+        rng : BRNGType, 
+        lambd : float = 1.0, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None, 
+    ) -> Tuple[BRNGType, BArrayType]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def random_normal(
+        self, 
+        shape: Union[int, Tuple[int, ...]], 
+        /,
+        *,
+        rng : BRNGType, 
+        mean : float = 0.0, std : float = 1.0, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None,
+    ) -> Tuple[BRNGType, BArrayType]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def random_geometric(
+        self, 
+        shape: Union[int, Tuple[int, ...]], 
+        /,
+        *, 
+        p : float, 
+        rng : BRNGType, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None,
+    ) -> Tuple[BRNGType, BArrayType]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def random_permutation(
+        self, 
+        n : int, 
+        /,
+        *, 
+        rng : BRNGType, 
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None
+    ) -> Tuple[BRNGType, BArrayType]:
+        raise NotImplementedError
+
+class ComputeBackend(ArrayAPINamespace[BArrayType, BDeviceType, BDtypeType], Protocol[BArrayType, BDeviceType, BDtypeType, BRNGType]):
+    simplified_name : str
+    ARRAY_TYPE : Type[BArrayType]
+    DEVICE_TYPE : Type[BDeviceType]
+    DTYPE_TYPE : Type[BDtypeType]
+    RNG_TYPE : Type[BRNGType]
+    default_integer_dtype : BDtypeType
+    default_floating_dtype : BDtypeType
+    default_boolean_dtype : BDtypeType
+    random : RNGBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
+
+    @abc.abstractmethod
+    def is_backendarray(self, data : Any) -> bool:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def from_numpy(
+        self, 
+        data : np.ndarray, 
+        /,
+        *,
+        dtype : Optional[BDtypeType] = None, 
+        device : Optional[BDeviceType] = None
+    ) -> BArrayType:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def from_other_backend(
+        self,
+        other_backend : "ComputeBackend",
+        data : ArrayAPIArray,
+        /
+    ) -> BArrayType:
+        """
+        Convert an array from another backend to this backend.
+        The other backend must be compatible with the ArrayAPI.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def to_numpy(
+        self, 
+        data : BArrayType
+    ) -> np.ndarray:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def to_dlpack(
+        self, 
+        data : BArrayType, 
+        /,
+    ) -> SupportsDLPack:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def dtype_is_real_integer(self, dtype : BDtypeType) -> bool:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def dtype_is_real_floating(self, dtype : BDtypeType) -> bool:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def dtype_is_boolean(self, dtype : BDtypeType) -> bool:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def abbreviate_array(self, x : BArrayType, try_cast_scalar: bool = True) -> Union[float, int, BArrayType]:
+        """
+        Abbreviates 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).
+        """
+        pass
+    
+    @abc.abstractmethod
+    def map_fn_over_arrays(self, data : Any, func : Callable[[BArrayType], BArrayType]) -> Any:
+        """
+        Map a function over arrays in a data structure and produce a new data structure with the same shape.
+        This is useful for applying a function to all arrays in a nested structure.
+        """

xbarray/common/implementations.py

@@ -0,0 +1,61 @@
+from typing import Any, Union, Callable
+from array_api_typing.typing_compat import ArrayAPINamespace as CompatNamespace, ArrayAPIArray as CompatArray, ArrayAPIDType as CompatDType
+
+__all__ = [
+    "get_abbreviate_array_function",
+    "get_map_fn_over_arrays_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=backend.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, dict):
+            return {k: map_fn_over_arrays(v, func) for k, v in data.items()}
+        elif isinstance(data, tuple):
+            return tuple(map_fn_over_arrays(i, func) for i in data)
+        elif isinstance(data, list):
+            return [map_fn_over_arrays(i, func) for i in data]
+        else:
+            return data
+    return map_fn_over_arrays

xbarray/jax/__init__.py

@@ -0,0 +1,25 @@
+import jax.numpy
+
+simplified_name = "jax"
+
+if hasattr(jax.numpy, "__array_api_version__"):
+    compat_module = jax.numpy
+    from jax.numpy import *
+else:
+    import jax.experimental.array_api as compat_module
+    from jax.experimental.array_api import *
+
+# 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
+    globals()[api_name] = partial(
+        getattr(array_api_extra, api_name),
+        xp=compat_module
+    )
+
+from ._typing import *
+from ._extra import *
+__import__(__package__ + ".random")

xbarray/jax/_extra.py

@@ -0,0 +1,98 @@
+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.base import ComputeBackend, SupportsDLPack
+
+__all__ = [
+    "default_integer_dtype",
+    "default_floating_dtype",
+    "default_boolean_dtype",
+    "is_backendarray",
+    "from_numpy",
+    "from_other_backend",
+    "to_numpy",
+    "to_dlpack",
+    "dtype_is_real_integer",
+    "dtype_is_real_floating",
+    "dtype_is_boolean",
+    "abbreviate_array",
+]
+
+default_integer_dtype = int
+default_floating_dtype = float
+default_boolean_dtype = bool
+
+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(
+    self,
+    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 xbarray.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
+    )

xbarray/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/jax/random.py

@@ -0,0 +1,116 @@
+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(0) 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
+
+@classmethod
+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/numpy/__init__.py

@@ -0,0 +1,19 @@
+from array_api_compat.numpy 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
+    globals()[api_name] = partial(
+        getattr(array_api_extra, api_name),
+        xp=compat_module
+    )
+
+from ._typing import *
+from ._extra import *
+__import__(__package__ + ".random")