array-api-compat 1.10.0__tar.gz → 1.11__tar.gz

{array_api_compat-1.10.0 → array_api_compat-1.11}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: array_api_compat
-Version: 1.10.0
+Version: 1.11
 Summary: A wrapper around NumPy and other array libraries to make them compatible with the Array API standard
 Home-page: https://data-apis.org/array-api-compat/
 Author: Consortium for Python Data API Standards
@@ -28,6 +28,15 @@ Provides-Extra: dask
 Requires-Dist: dask; extra == "dask"
 Provides-Extra: sparse
 Requires-Dist: sparse>=0.15.1; extra == "sparse"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-python
+Dynamic: summary
 
 # Array API compatibility library
 

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/__init__.py

@@ -17,6 +17,6 @@ to ensure they are not using functionality outside of the standard, but prefer
 this implementation for the default when working with NumPy arrays.
 
 """
-__version__ = '1.10.0'
+__version__ = '1.11'
 
 from .common import *  # noqa: F401, F403

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/common/_aliases.py

@@ -233,11 +233,6 @@ def unique_values(x: ndarray, /, xp) -> ndarray:
         **kwargs,
     )
 
-def astype(x: ndarray, dtype: Dtype, /, *, copy: bool = True) -> ndarray:
-    if not copy and dtype == x.dtype:
-        return x
-    return x.astype(dtype=dtype, copy=copy)
-
 # These functions have different keyword argument names
 
 def std(
@@ -297,6 +292,36 @@ def cumulative_sum(
         )
     return res
 
+
+def cumulative_prod(
+    x: ndarray,
+    /,
+    xp,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[Dtype] = None,
+    include_initial: bool = False,
+    **kwargs
+) -> ndarray:
+    wrapped_xp = array_namespace(x)
+
+    if axis is None:
+        if x.ndim > 1:
+            raise ValueError("axis must be specified in cumulative_prod for more than one dimension")
+        axis = 0
+
+    res = xp.cumprod(x, axis=axis, dtype=dtype, **kwargs)
+
+    # np.cumprod does not support include_initial
+    if include_initial:
+        initial_shape = list(x.shape)
+        initial_shape[axis] = 1
+        res = xp.concatenate(
+            [wrapped_xp.ones(shape=initial_shape, dtype=res.dtype, device=device(res)), res],
+            axis=axis,
+        )
+    return res
+
 # The min and max argument names in clip are different and not optional in numpy, and type
 # promotion behavior is different.
 def clip(
@@ -549,7 +574,7 @@ __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
-           'astype', 'std', 'var', 'cumulative_sum', 'clip', 'permute_dims',
+           'std', 'var', 'cumulative_sum', 'cumulative_prod','clip', 'permute_dims',
            'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
            'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
            'unstack', 'sign']

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/common/_fft.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Union, Optional, Literal
 
 if TYPE_CHECKING:
-    from ._typing import Device, ndarray
+    from ._typing import Device, ndarray, DType
     from collections.abc import Sequence
 
 # Note: NumPy fft functions improperly upcast float32 and complex64 to
@@ -149,15 +149,37 @@ def ihfft(
         return res.astype(xp.complex64)
     return res
 
-def fftfreq(n: int, /, xp, *, d: float = 1.0, device: Optional[Device] = None) -> ndarray:
+def fftfreq(
+    n: int,
+    /,
+    xp,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> ndarray:
     if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
-    return xp.fft.fftfreq(n, d=d)
+    res = xp.fft.fftfreq(n, d=d)
+    if dtype is not None:
+        return res.astype(dtype)
+    return res
 
-def rfftfreq(n: int, /, xp, *, d: float = 1.0, device: Optional[Device] = None) -> ndarray:
+def rfftfreq(
+    n: int,
+    /,
+    xp,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> ndarray:
     if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
-    return xp.fft.rfftfreq(n, d=d)
+    res = xp.fft.rfftfreq(n, d=d)
+    if dtype is not None:
+        return res.astype(dtype)
+    return res
 
 def fftshift(x: ndarray, /, xp, *, axes: Union[int, Sequence[int]] = None) -> ndarray:
     return xp.fft.fftshift(x, axes=axes)

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/common/_helpers.py

@@ -11,14 +11,14 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Optional, Union, Any
-    from ._typing import Array, Device
+    from ._typing import Array, Device, Namespace
 
 import sys
 import math
 import inspect
 import warnings
 
-def _is_jax_zero_gradient_array(x):
+def _is_jax_zero_gradient_array(x: object) -> bool:
     """Return True if `x` is a zero-gradient array.
 
     These arrays are a design quirk of Jax that may one day be removed.
@@ -32,7 +32,8 @@ def _is_jax_zero_gradient_array(x):
 
     return isinstance(x, np.ndarray) and x.dtype == jax.float0
 
-def is_numpy_array(x):
+
+def is_numpy_array(x: object) -> bool:
     """
     Return True if `x` is a NumPy array.
 
@@ -63,7 +64,8 @@ def is_numpy_array(x):
     return (isinstance(x, (np.ndarray, np.generic))
             and not _is_jax_zero_gradient_array(x))
 
-def is_cupy_array(x):
+
+def is_cupy_array(x: object) -> bool:
     """
     Return True if `x` is a CuPy array.
 
@@ -93,7 +95,8 @@ def is_cupy_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, cp.ndarray)
 
-def is_torch_array(x):
+
+def is_torch_array(x: object) -> bool:
     """
     Return True if `x` is a PyTorch tensor.
 
@@ -120,7 +123,8 @@ def is_torch_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, torch.Tensor)
 
-def is_ndonnx_array(x):
+
+def is_ndonnx_array(x: object) -> bool:
     """
     Return True if `x` is a ndonnx Array.
 
@@ -147,7 +151,8 @@ def is_ndonnx_array(x):
 
     return isinstance(x, ndx.Array)
 
-def is_dask_array(x):
+
+def is_dask_array(x: object) -> bool:
     """
     Return True if `x` is a dask.array Array.
 
@@ -174,7 +179,8 @@ def is_dask_array(x):
 
     return isinstance(x, dask.array.Array)
 
-def is_jax_array(x):
+
+def is_jax_array(x: object) -> bool:
     """
     Return True if `x` is a JAX array.
 
@@ -202,6 +208,7 @@ def is_jax_array(x):
 
     return isinstance(x, jax.Array) or _is_jax_zero_gradient_array(x)
 
+
 def is_pydata_sparse_array(x) -> bool:
     """
     Return True if `x` is an array from the `sparse` package.
@@ -231,7 +238,8 @@ def is_pydata_sparse_array(x) -> bool:
     # TODO: Account for other backends.
     return isinstance(x, sparse.SparseArray)
 
-def is_array_api_obj(x):
+
+def is_array_api_obj(x: object) -> bool:
     """
     Return True if `x` is an array API compatible array object.
 
@@ -254,10 +262,12 @@ def is_array_api_obj(x):
         or is_pydata_sparse_array(x) \
         or hasattr(x, '__array_namespace__')
 
-def _compat_module_name():
+
+def _compat_module_name() -> str:
     assert __name__.endswith('.common._helpers')
     return __name__.removesuffix('.common._helpers')
 
+
 def is_numpy_namespace(xp) -> bool:
     """
     Returns True if `xp` is a NumPy namespace.
@@ -278,6 +288,7 @@ def is_numpy_namespace(xp) -> bool:
     """
     return xp.__name__ in {'numpy', _compat_module_name() + '.numpy'}
 
+
 def is_cupy_namespace(xp) -> bool:
     """
     Returns True if `xp` is a CuPy namespace.
@@ -298,6 +309,7 @@ def is_cupy_namespace(xp) -> bool:
     """
     return xp.__name__ in {'cupy', _compat_module_name() + '.cupy'}
 
+
 def is_torch_namespace(xp) -> bool:
     """
     Returns True if `xp` is a PyTorch namespace.
@@ -319,7 +331,7 @@ def is_torch_namespace(xp) -> bool:
     return xp.__name__ in {'torch', _compat_module_name() + '.torch'}
 
 
-def is_ndonnx_namespace(xp):
+def is_ndonnx_namespace(xp) -> bool:
     """
     Returns True if `xp` is an NDONNX namespace.
 
@@ -337,7 +349,8 @@ def is_ndonnx_namespace(xp):
     """
     return xp.__name__ == 'ndonnx'
 
-def is_dask_namespace(xp):
+
+def is_dask_namespace(xp) -> bool:
     """
     Returns True if `xp` is a Dask namespace.
 
@@ -357,7 +370,8 @@ def is_dask_namespace(xp):
     """
     return xp.__name__ in {'dask.array', _compat_module_name() + '.dask.array'}
 
-def is_jax_namespace(xp):
+
+def is_jax_namespace(xp) -> bool:
     """
     Returns True if `xp` is a JAX namespace.
 
@@ -378,7 +392,8 @@ def is_jax_namespace(xp):
     """
     return xp.__name__ in {'jax.numpy', 'jax.experimental.array_api'}
 
-def is_pydata_sparse_namespace(xp):
+
+def is_pydata_sparse_namespace(xp) -> bool:
     """
     Returns True if `xp` is a pydata/sparse namespace.
 
@@ -396,7 +411,8 @@ def is_pydata_sparse_namespace(xp):
     """
     return xp.__name__ == 'sparse'
 
-def is_array_api_strict_namespace(xp):
+
+def is_array_api_strict_namespace(xp) -> bool:
     """
     Returns True if `xp` is an array-api-strict namespace.
 
@@ -414,14 +430,16 @@ def is_array_api_strict_namespace(xp):
     """
     return xp.__name__ == 'array_api_strict'
 
-def _check_api_version(api_version):
-    if api_version in ['2021.12', '2022.12']:
-        warnings.warn(f"The {api_version} version of the array API specification was requested but the returned namespace is actually version 2023.12")
+
+def _check_api_version(api_version: str) -> None:
+    if api_version in ['2021.12', '2022.12', '2023.12']:
+        warnings.warn(f"The {api_version} version of the array API specification was requested but the returned namespace is actually version 2024.12")
     elif api_version is not None and api_version not in ['2021.12', '2022.12',
-                                                         '2023.12']:
-        raise ValueError("Only the 2023.12 version of the array API specification is currently supported")
+                                                         '2023.12', '2024.12']:
+        raise ValueError("Only the 2024.12 version of the array API specification is currently supported")
 
-def array_namespace(*xs, api_version=None, use_compat=None):
+
+def array_namespace(*xs, api_version=None, use_compat=None) -> Namespace:
     """
     Get the array API compatible namespace for the arrays `xs`.
 
@@ -433,7 +451,7 @@ def array_namespace(*xs, api_version=None, use_compat=None):
 
     api_version: str
         The newest version of the spec that you need support for (currently
-        the compat library wrapped APIs support v2023.12).
+        the compat library wrapped APIs support v2024.12).
 
     use_compat: bool or None
         If None (the default), the native namespace will be returned if it is
@@ -630,24 +648,24 @@ def device(x: Array, /) -> Device:
     if is_numpy_array(x):
         return "cpu"
     elif is_dask_array(x):
-        # Peek at the metadata of the jax array to determine type
-        try:
-            import numpy as np
-            if isinstance(x._meta, np.ndarray):
-                # Must be on CPU since backed by numpy
-                return "cpu"
-        except ImportError:
-            pass
+        # Peek at the metadata of the Dask array to determine type
+        if is_numpy_array(x._meta):
+            # Must be on CPU since backed by numpy
+            return "cpu"
         return _DASK_DEVICE
     elif is_jax_array(x):
-        # JAX has .device() as a method, but it is being deprecated so that it
-        # can become a property, in accordance with the standard. In order for
-        # this function to not break when JAX makes the flip, we check for
-        # both here.
-        if inspect.ismethod(x.device):
-            return x.device()
+        # FIXME Jitted JAX arrays do not have a device attribute
+        #       https://github.com/jax-ml/jax/issues/26000
+        #       Return None in this case. Note that this workaround breaks
+        #       the standard and will result in new arrays being created on the
+        #       default device instead of the same device as the input array(s).
+        x_device = getattr(x, 'device', None)
+        # Older JAX releases had .device() as a method, which has been replaced
+        # with a property in accordance with the standard.
+        if inspect.ismethod(x_device):
+            return x_device()
         else:
-            return x.device
+            return x_device
     elif is_pydata_sparse_array(x):
         # `sparse` will gain `.device`, so check for this first.
         x_device = getattr(x, 'device', None)
@@ -778,8 +796,11 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
         raise ValueError(f"Unsupported device {device!r}")
     elif is_jax_array(x):
         if not hasattr(x, "__array_namespace__"):
-            # In JAX v0.4.31 and older, this import adds to_device method to x.
+            # In JAX v0.4.31 and older, this import adds to_device method to x...
             import jax.experimental.array_api # noqa: F401
+            # ... but only on eager JAX. It won't work inside jax.jit.
+            if not hasattr(x, "to_device"):
+                return x
         return x.to_device(device, stream=stream)
     elif is_pydata_sparse_array(x) and device == _device(x):
         # Perform trivial check to return the same array if
@@ -788,24 +809,30 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
     return x.to_device(device, stream=stream)
 
 
-def size(x):
+def size(x: Array) -> int | None:
     """
     Return the total number of elements of x.
 
     This is equivalent to `x.size` according to the `standard
     <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.size.html>`__.
+
     This helper is included because PyTorch defines `size` in an
     :external+torch:meth:`incompatible way <torch.Tensor.size>`.
-
+    It also fixes dask.array's behaviour which returns nan for unknown sizes, whereas
+    the standard requires None.
     """
+    # Lazy API compliant arrays, such as ndonnx, can contain None in their shape
     if None in x.shape:
         return None
-    return math.prod(x.shape)
+    out = math.prod(x.shape)
+    # dask.array.Array.shape can contain NaN
+    return None if math.isnan(out) else out
 
 
-def is_writeable_array(x) -> bool:
+def is_writeable_array(x: object) -> bool:
     """
     Return False if ``x.__setitem__`` is expected to raise; True otherwise.
+    Return False if `x` is not an array API compatible object.
 
     Warning
     -------
@@ -816,7 +843,67 @@ def is_writeable_array(x) -> bool:
         return x.flags.writeable
     if is_jax_array(x) or is_pydata_sparse_array(x):
         return False
-    return True
+    return is_array_api_obj(x)
+
+
+def is_lazy_array(x: object) -> bool:
+    """Return True if x is potentially a future or it may be otherwise impossible or
+    expensive to eagerly read its contents, regardless of their size, e.g. by
+    calling ``bool(x)`` or ``float(x)``.
+
+    Return False otherwise; e.g. ``bool(x)`` etc. is guaranteed to succeed and to be
+    cheap as long as the array has the right dtype and size.
+
+    Note
+    ----
+    This function errs on the side of caution for array types that may or may not be
+    lazy, e.g. JAX arrays, by always returning True for them.
+    """
+    if (
+        is_numpy_array(x)
+        or is_cupy_array(x)
+        or is_torch_array(x)
+        or is_pydata_sparse_array(x)
+    ):
+        return False
+
+    # **JAX note:** while it is possible to determine if you're inside or outside
+    # jax.jit by testing the subclass of a jax.Array object, as well as testing bool()
+    # as we do below for unknown arrays, this is not recommended by JAX best practices.
+
+    # **Dask note:** Dask eagerly computes the graph on __bool__, __float__, and so on.
+    # This behaviour, while impossible to change without breaking backwards
+    # compatibility, is highly detrimental to performance as the whole graph will end
+    # up being computed multiple times.
+
+    if is_jax_array(x) or is_dask_array(x) or is_ndonnx_array(x):
+        return True
+
+    if not is_array_api_obj(x):
+        return False
+
+    # Unknown Array API compatible object. Note that this test may have dire consequences
+    # in terms of performance, e.g. for a lazy object that eagerly computes the graph
+    # on __bool__ (dask is one such example, which however is special-cased above).
+
+    # Select a single point of the array
+    s = size(x)
+    if s is None:
+        return True
+    xp = array_namespace(x)
+    if s > 1:
+        x = xp.reshape(x, (-1,))[0]
+    # Cast to dtype=bool and deal with size 0 arrays
+    x = xp.any(x)
+
+    try:
+        bool(x)
+        return False
+    # The Array API standard dictactes that __bool__ should raise TypeError if the
+    # output cannot be defined.
+    # Here we allow for it to raise arbitrary exceptions, e.g. like Dask does.
+    except Exception:
+        return True
 
 
 __all__ = [
@@ -840,6 +927,7 @@ __all__ = [
     "is_pydata_sparse_array",
     "is_pydata_sparse_namespace",
     "is_writeable_array",
+    "is_lazy_array",
     "size",
     "to_device",
 ]

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/common/_typing.py

@@ -5,6 +5,7 @@ __all__ = [
     "SupportsBufferProtocol",
 ]
 
+from types import ModuleType
 from typing import (
     Any,
     TypeVar,
@@ -21,3 +22,5 @@ SupportsBufferProtocol = Any
 
 Array = Any
 Device = Any
+DType = Any
+Namespace = ModuleType

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/cupy/__init__.py

@@ -13,4 +13,4 @@ __import__(__package__ + '.fft')
 
 from ..common._helpers import * # noqa: F401,F403
 
-__array_api_version__ = '2023.12'
+__array_api_version__ = '2024.12'

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/cupy/_aliases.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import cupy as cp
 
-from ..common import _aliases
+from ..common import _aliases, _helpers
 from .._internal import get_xp
 
 from ._info import __array_namespace_info__
@@ -46,10 +46,10 @@ unique_all = get_xp(cp)(_aliases.unique_all)
 unique_counts = get_xp(cp)(_aliases.unique_counts)
 unique_inverse = get_xp(cp)(_aliases.unique_inverse)
 unique_values = get_xp(cp)(_aliases.unique_values)
-astype = _aliases.astype
 std = get_xp(cp)(_aliases.std)
 var = get_xp(cp)(_aliases.var)
 cumulative_sum = get_xp(cp)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(cp)(_aliases.cumulative_prod)
 clip = get_xp(cp)(_aliases.clip)
 permute_dims = get_xp(cp)(_aliases.permute_dims)
 reshape = get_xp(cp)(_aliases.reshape)
@@ -110,6 +110,21 @@ def asarray(
 
         return cp.array(obj, dtype=dtype, **kwargs)
 
+
+def astype(
+    x: ndarray,
+    dtype: Dtype,
+    /,
+    *,
+    copy: bool = True,
+    device: Optional[Device] = None,
+) -> ndarray:
+    if device is None:
+        return x.astype(dtype=dtype, copy=copy)
+    out = _helpers.to_device(x.astype(dtype=dtype, copy=False), device)
+    return out.copy() if copy and out is x else out
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(cp, 'vecdot'):
@@ -127,10 +142,10 @@ if hasattr(cp, 'unstack'):
 else:
     unstack = get_xp(cp)(_aliases.unstack)
 
-__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'bool',
+__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'astype',
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'concat', 'pow', 'sign']
+                              'bool', 'concat', 'pow', 'sign']
 
 _all_ignore = ['cp', 'get_xp']

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/cupy/_info.py

@@ -101,7 +101,7 @@ class __array_namespace_info__:
             "boolean indexing": True,
             "data-dependent shapes": True,
             # 'max rank' will be part of the 2024.12 standard
-            # "max rank": 64,
+            "max dimensions": 64,
         }
 
     def default_device(self):

{array_api_compat-1.10.0 → array_api_compat-1.11}/array_api_compat/dask/array/__init__.py

@@ -3,7 +3,7 @@ from dask.array import * # noqa: F403
 # These imports may overwrite names from the import * above.
 from ._aliases import * # noqa: F403
 
-__array_api_version__ = '2023.12'
+__array_api_version__ = '2024.12'
 
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')