array-api-compat 1.8__tar.gz → 1.9__tar.gz

{array_api_compat-1.8 → array_api_compat-1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: array_api_compat
-Version: 1.8
+Version: 1.9
 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

{array_api_compat-1.8 → array_api_compat-1.9}/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.8'
+__version__ = '1.9'
 
 from .common import *  # noqa: F401, F403

{array_api_compat-1.8 → array_api_compat-1.9}/array_api_compat/common/_aliases.py

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
 from typing import NamedTuple
 import inspect
 
-from ._helpers import array_namespace, _check_device
+from ._helpers import array_namespace, _check_device, device, is_torch_array
 
 # These functions are modified from the NumPy versions.
 
@@ -264,6 +264,38 @@ def var(
 ) -> ndarray:
     return xp.var(x, axis=axis, ddof=correction, keepdims=keepdims, **kwargs)
 
+# cumulative_sum is renamed from cumsum, and adds the include_initial keyword
+# argument
+
+def cumulative_sum(
+    x: ndarray,
+    /,
+    xp,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[Dtype] = None,
+    include_initial: bool = False,
+    **kwargs
+) -> ndarray:
+    wrapped_xp = array_namespace(x)
+
+    # TODO: The standard is not clear about what should happen when x.ndim == 0.
+    if axis is None:
+        if x.ndim > 1:
+            raise ValueError("axis must be specified in cumulative_sum for more than one dimension")
+        axis = 0
+
+    res = xp.cumsum(x, axis=axis, dtype=dtype, **kwargs)
+
+    # np.cumsum does not support include_initial
+    if include_initial:
+        initial_shape = list(x.shape)
+        initial_shape[axis] = 1
+        res = xp.concatenate(
+            [wrapped_xp.zeros(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.
@@ -281,10 +313,11 @@ def clip(
         return isinstance(a, (int, float, type(None)))
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
-    result_shape = xp.broadcast_shapes(x.shape, min_shape, max_shape)
 
     wrapped_xp = array_namespace(x)
 
+    result_shape = xp.broadcast_shapes(x.shape, min_shape, max_shape)
+
     # np.clip does type promotion but the array API clip requires that the
     # output have the same dtype as x. We do this instead of just downcasting
     # the result of xp.clip() to handle some corner cases better (e.g.,
@@ -305,20 +338,26 @@ def clip(
 
     # At least handle the case of Python integers correctly (see
     # https://github.com/numpy/numpy/pull/26892).
-    if type(min) is int and min <= xp.iinfo(x.dtype).min:
+    if type(min) is int and min <= wrapped_xp.iinfo(x.dtype).min:
         min = None
-    if type(max) is int and max >= xp.iinfo(x.dtype).max:
+    if type(max) is int and max >= wrapped_xp.iinfo(x.dtype).max:
         max = None
 
     if out is None:
-        out = wrapped_xp.asarray(xp.broadcast_to(x, result_shape), copy=True)
+        out = wrapped_xp.asarray(xp.broadcast_to(x, result_shape),
+                                 copy=True, device=device(x))
     if min is not None:
-        a = xp.broadcast_to(xp.asarray(min), result_shape)
+        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(min):
+            # Avoid loss of precision due to torch defaulting to float32
+            min = wrapped_xp.asarray(min, dtype=xp.float64)
+        a = xp.broadcast_to(wrapped_xp.asarray(min, device=device(x)), result_shape)
         ia = (out < a) | xp.isnan(a)
         # torch requires an explicit cast here
         out[ia] = wrapped_xp.astype(a[ia], out.dtype)
     if max is not None:
-        b = xp.broadcast_to(xp.asarray(max), result_shape)
+        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(max):
+            max = wrapped_xp.asarray(max, dtype=xp.float64)
+        b = xp.broadcast_to(wrapped_xp.asarray(max, device=device(x)), result_shape)
         ib = (out > b) | xp.isnan(b)
         out[ib] = wrapped_xp.astype(b[ib], out.dtype)
     # Return a scalar for 0-D
@@ -389,42 +428,6 @@ def nonzero(x: ndarray, /, xp, **kwargs) -> Tuple[ndarray, ...]:
         raise ValueError("nonzero() does not support zero-dimensional arrays")
     return xp.nonzero(x, **kwargs)
 
-# sum() and prod() should always upcast when dtype=None
-def sum(
-    x: ndarray,
-    /,
-    xp,
-    *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
-    keepdims: bool = False,
-    **kwargs,
-) -> ndarray:
-    # `xp.sum` already upcasts integers, but not floats or complexes
-    if dtype is None:
-        if x.dtype == xp.float32:
-            dtype = xp.float64
-        elif x.dtype == xp.complex64:
-            dtype = xp.complex128
-    return xp.sum(x, axis=axis, dtype=dtype, keepdims=keepdims, **kwargs)
-
-def prod(
-    x: ndarray,
-    /,
-    xp,
-    *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
-    keepdims: bool = False,
-    **kwargs,
-) -> ndarray:
-    if dtype is None:
-        if x.dtype == xp.float32:
-            dtype = xp.float64
-        elif x.dtype == xp.complex64:
-            dtype = xp.complex128
-    return xp.prod(x, dtype=dtype, axis=axis, keepdims=keepdims, **kwargs)
-
 # ceil, floor, and trunc return integers for integer inputs
 
 def ceil(x: ndarray, /, xp, **kwargs) -> ndarray:
@@ -521,10 +524,17 @@ def isdtype(
         # array_api_strict implementation will be very strict.
         return dtype == kind
 
+# unstack is a new function in the 2023.12 array API standard
+def unstack(x: ndarray, /, xp, *, axis: int = 0) -> Tuple[ndarray, ...]:
+    if x.ndim == 0:
+        raise ValueError("Input array must be at least 1-d.")
+    return tuple(xp.moveaxis(x, axis, 0))
+
 __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', 'clip', 'permute_dims', 'reshape', 'argsort',
-           'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
-           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']
+           'astype', 'std', 'var', 'cumulative_sum', 'clip', 'permute_dims',
+           'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
+           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
+           'unstack']

{array_api_compat-1.8 → array_api_compat-1.9}/array_api_compat/common/_helpers.py

@@ -145,7 +145,7 @@ def is_ndonnx_array(x):
 
     import ndonnx as ndx
 
-    return isinstance(x, ndx.Array)  
+    return isinstance(x, ndx.Array)
 
 def is_dask_array(x):
     """
@@ -202,7 +202,6 @@ 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.
@@ -255,6 +254,166 @@ def is_array_api_obj(x):
         or is_pydata_sparse_array(x) \
         or hasattr(x, '__array_namespace__')
 
+def _compat_module_name():
+    assert __name__.endswith('.common._helpers')
+    return __name__.removesuffix('.common._helpers')
+
+def is_numpy_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a NumPy namespace.
+
+    This includes both NumPy itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'numpy', _compat_module_name() + '.numpy'}
+
+def is_cupy_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a CuPy namespace.
+
+    This includes both CuPy itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'cupy', _compat_module_name() + '.cupy'}
+
+def is_torch_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a PyTorch namespace.
+
+    This includes both PyTorch itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'torch', _compat_module_name() + '.torch'}
+    
+
+def is_ndonnx_namespace(xp):
+    """
+    Returns True if `xp` is an NDONNX namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ == 'ndonnx'
+
+def is_dask_namespace(xp):
+    """
+    Returns True if `xp` is a Dask namespace.
+
+    This includes both ``dask.array`` itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'dask.array', _compat_module_name() + '.dask.array'}
+
+def is_jax_namespace(xp):
+    """
+    Returns True if `xp` is a JAX namespace.
+
+    This includes ``jax.numpy`` and ``jax.experimental.array_api`` which existed in
+    older versions of JAX.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'jax.numpy', 'jax.experimental.array_api'}
+
+def is_pydata_sparse_namespace(xp):
+    """
+    Returns True if `xp` is a pydata/sparse namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ == 'sparse'
+
+def is_array_api_strict_namespace(xp):
+    """
+    Returns True if `xp` is an array-api-strict namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    """
+    return xp.__name__ == 'array_api_strict'
+
 def _check_api_version(api_version):
     if api_version == '2021.12':
         warnings.warn("The 2021.12 version of the array API specification was requested but the returned namespace is actually version 2022.12")
@@ -340,12 +499,9 @@ def array_namespace(*xs, api_version=None, use_compat=None):
             elif use_compat is False:
                 namespaces.add(np)
             else:
-                # numpy 2.0 has __array_namespace__ and is fully array API
+                # numpy 2.0+ have __array_namespace__, however, they are not yet fully array API
                 # compatible.
-                if hasattr(x, '__array_namespace__'):
-                    namespaces.add(x.__array_namespace__(api_version=api_version))
-                else:
-                    namespaces.add(numpy_namespace)
+                namespaces.add(numpy_namespace)
         elif is_cupy_array(x):
             if _use_compat:
                 _check_api_version(api_version)
@@ -377,9 +533,13 @@ def array_namespace(*xs, api_version=None, use_compat=None):
             elif use_compat is False:
                 import jax.numpy as jnp
             else:
-                # jax.experimental.array_api is already an array namespace. We do
-                # not have a wrapper submodule for it.
-                import jax.experimental.array_api as jnp
+                # JAX v0.4.32 and newer implements the array API directly in jax.numpy.
+                # For older JAX versions, it is available via jax.experimental.array_api.
+                import jax.numpy
+                if hasattr(jax.numpy, "__array_api_version__"):
+                    jnp = jax.numpy
+                else:
+                    import jax.experimental.array_api as jnp
             namespaces.add(jnp)
         elif is_pydata_sparse_array(x):
             if use_compat is True:
@@ -613,8 +773,9 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
             return x
         raise ValueError(f"Unsupported device {device!r}")
     elif is_jax_array(x):
-        # This import adds to_device to x
-        import jax.experimental.array_api # noqa: F401
+        if not hasattr(x, "__array_namespace__"):
+            # In JAX v0.4.31 and older, this import adds to_device method to x.
+            import jax.experimental.array_api # noqa: F401
         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
@@ -641,13 +802,21 @@ __all__ = [
     "device",
     "get_namespace",
     "is_array_api_obj",
+    "is_array_api_strict_namespace",
     "is_cupy_array",
+    "is_cupy_namespace",
     "is_dask_array",
+    "is_dask_namespace",
     "is_jax_array",
+    "is_jax_namespace",
     "is_numpy_array",
+    "is_numpy_namespace",
     "is_torch_array",
+    "is_torch_namespace",
     "is_ndonnx_array",
+    "is_ndonnx_namespace",
     "is_pydata_sparse_array",
+    "is_pydata_sparse_namespace",
     "size",
     "to_device",
 ]

{array_api_compat-1.8 → array_api_compat-1.9}/array_api_compat/common/_linalg.py

@@ -147,11 +147,6 @@ def diagonal(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray:
     return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs)
 
 def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarray:
-    if dtype is None:
-        if x.dtype == xp.float32:
-            dtype = xp.float64
-        elif x.dtype == xp.complex64:
-            dtype = xp.complex128
     return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))
 
 __all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',

{array_api_compat-1.8 → array_api_compat-1.9}/array_api_compat/cupy/_aliases.py

@@ -5,6 +5,8 @@ import cupy as cp
 from ..common import _aliases
 from .._internal import get_xp
 
+from ._info import __array_namespace_info__
+
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from typing import Optional, Union
@@ -47,14 +49,13 @@ 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)
 clip = get_xp(cp)(_aliases.clip)
 permute_dims = get_xp(cp)(_aliases.permute_dims)
 reshape = get_xp(cp)(_aliases.reshape)
 argsort = get_xp(cp)(_aliases.argsort)
 sort = get_xp(cp)(_aliases.sort)
 nonzero = get_xp(cp)(_aliases.nonzero)
-sum = get_xp(cp)(_aliases.sum)
-prod = get_xp(cp)(_aliases.prod)
 ceil = get_xp(cp)(_aliases.ceil)
 floor = get_xp(cp)(_aliases.floor)
 trunc = get_xp(cp)(_aliases.trunc)
@@ -108,20 +109,34 @@ def asarray(
 
         return cp.array(obj, dtype=dtype, **kwargs)
 
+def sign(x: ndarray, /) -> ndarray:
+    # CuPy sign() does not propagate nans. See
+    # https://github.com/data-apis/array-api-compat/issues/136
+    out = cp.sign(x)
+    out[cp.isnan(x)] = cp.nan
+    return 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'):
     vecdot = cp.vecdot
 else:
     vecdot = get_xp(cp)(_aliases.vecdot)
+
 if hasattr(cp, 'isdtype'):
     isdtype = cp.isdtype
 else:
     isdtype = get_xp(cp)(_aliases.isdtype)
 
-__all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
-                              'acosh', 'asin', 'asinh', 'atan', 'atan2',
-                              'atanh', 'bitwise_left_shift', 'bitwise_invert',
-                              'bitwise_right_shift', 'concat', 'pow']
+if hasattr(cp, 'unstack'):
+    unstack = cp.unstack
+else:
+    unstack = get_xp(cp)(_aliases.unstack)
+
+__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'bool',
+                              'acos', 'acosh', 'asin', 'asinh', 'atan',
+                              'atan2', 'atanh', 'bitwise_left_shift',
+                              'bitwise_invert', 'bitwise_right_shift',
+                              'concat', 'pow', 'sign']
 
 _all_ignore = ['cp', 'get_xp']