array-api-compat 1.5__tar.gz → 1.6__tar.gz

array_api_compat-1.6/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.1
+Name: array_api_compat
+Version: 1.6
+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
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: numpy
+Requires-Dist: numpy; extra == "numpy"
+Provides-Extra: cupy
+Requires-Dist: cupy; extra == "cupy"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Provides-Extra: pytorch
+Requires-Dist: pytorch; extra == "pytorch"
+Provides-Extra: dask
+Requires-Dist: dask; extra == "dask"
+
+# Array API compatibility library
+
+This is a small wrapper around common array libraries that is compatible with
+the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
+NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
+libraries, or if you encounter any issues, please [open an
+issue](https://github.com/data-apis/array-api-compat/issues).
+
+See the documentation for more details https://data-apis.org/array-api-compat/

array_api_compat-1.6/README.md

@@ -0,0 +1,9 @@
+# Array API compatibility library
+
+This is a small wrapper around common array libraries that is compatible with
+the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
+NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
+libraries, or if you encounter any issues, please [open an
+issue](https://github.com/data-apis/array-api-compat/issues).
+
+See the documentation for more details https://data-apis.org/array-api-compat/

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

{array_api_compat-1.5 → array_api_compat-1.6}/array_api_compat/common/_aliases.py

@@ -6,18 +6,18 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    import numpy as np
     from typing import Optional, Sequence, Tuple, Union
-    from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
+    from ._typing import ndarray, Device, Dtype
 
 from typing import NamedTuple
-from types import ModuleType
 import inspect
 
-from ._helpers import _check_device, is_numpy_array, array_namespace
+from ._helpers import _check_device
 
 # These functions are modified from the NumPy versions.
 
+# Creation functions add the device keyword (which does nothing for NumPy)
+
 def arange(
     start: Union[int, float],
     /,
@@ -268,91 +268,6 @@ def var(
 def permute_dims(x: ndarray, /, axes: Tuple[int, ...], xp) -> ndarray:
     return xp.transpose(x, axes)
 
-# Creation functions add the device keyword (which does nothing for NumPy)
-
-# asarray also adds the copy keyword
-def _asarray(
-    obj: Union[
-        ndarray,
-        bool,
-        int,
-        float,
-        NestedSequence[bool | int | float],
-        SupportsBufferProtocol,
-    ],
-    /,
-    *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
-    copy: "Optional[Union[bool, np._CopyMode]]" = None,
-    namespace = None,
-    **kwargs,
-) -> ndarray:
-    """
-    Array API compatibility wrapper for asarray().
-
-    See the corresponding documentation in NumPy/CuPy and/or the array API
-    specification for more details.
-
-    """
-    if namespace is None:
-        try:
-            xp = array_namespace(obj, _use_compat=False)
-        except ValueError:
-            # TODO: What about lists of arrays?
-            raise ValueError("A namespace must be specified for asarray() with non-array input")
-    elif isinstance(namespace, ModuleType):
-        xp = namespace
-    elif namespace == 'numpy':
-        import numpy as xp
-    elif namespace == 'cupy':
-        import cupy as xp
-    elif namespace == 'dask.array':
-        import dask.array as xp
-    else:
-        raise ValueError("Unrecognized namespace argument to asarray()")
-
-    _check_device(xp, device)
-    if is_numpy_array(obj):
-        import numpy as np
-        if hasattr(np, '_CopyMode'):
-            # Not present in older NumPys
-            COPY_FALSE = (False, np._CopyMode.IF_NEEDED)
-            COPY_TRUE = (True, np._CopyMode.ALWAYS)
-        else:
-            COPY_FALSE = (False,)
-            COPY_TRUE = (True,)
-    else:
-        COPY_FALSE = (False,)
-        COPY_TRUE = (True,)
-    if copy in COPY_FALSE:
-        # copy=False is not yet implemented in xp.asarray
-        raise NotImplementedError("copy=False is not yet implemented")
-    if (hasattr(xp, "ndarray") and isinstance(obj, xp.ndarray)) or hasattr(obj, "__array__"):
-        #print('hit me')
-        if dtype is not None and obj.dtype != dtype:
-            copy = True
-        #print(copy)
-        if copy in COPY_TRUE:
-            copy_kwargs = {}
-            if namespace != "dask.array":
-                copy_kwargs["copy"] = True
-            else:
-                # No copy kw in dask.asarray so we go thorugh np.asarray first
-                # (like dask also does) but copy after
-                if dtype is None:
-                    # Same dtype copy is no-op in dask
-                    #print("in here?")
-                    return obj.copy()
-                import numpy as np
-                #print(obj)
-                obj = np.asarray(obj).copy()
-                #print(obj)
-            return xp.array(obj, dtype=dtype, **copy_kwargs)
-        return obj
-
-    return xp.asarray(obj, dtype=dtype, **kwargs)
-
 # np.reshape calls the keyword argument 'newshape' instead of 'shape'
 def reshape(x: ndarray,
             /,

{array_api_compat-1.5 → array_api_compat-1.6}/array_api_compat/common/_helpers.py

@@ -19,6 +19,24 @@ import inspect
 import warnings
 
 def is_numpy_array(x):
+    """
+    Return True if `x` is a NumPy array.
+
+    This function does not import NumPy if it has not already been imported
+    and is therefore cheap to use.
+
+    This also returns True for `ndarray` subclasses and NumPy scalar objects.
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
+    """
     # Avoid importing NumPy if it isn't already
     if 'numpy' not in sys.modules:
         return False
@@ -29,6 +47,24 @@ def is_numpy_array(x):
     return isinstance(x, (np.ndarray, np.generic))
 
 def is_cupy_array(x):
+    """
+    Return True if `x` is a CuPy array.
+
+    This function does not import CuPy if it has not already been imported
+    and is therefore cheap to use.
+
+    This also returns True for `cupy.ndarray` subclasses and CuPy scalar objects.
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
+    """
     # Avoid importing NumPy if it isn't already
     if 'cupy' not in sys.modules:
         return False
@@ -39,6 +75,22 @@ def is_cupy_array(x):
     return isinstance(x, (cp.ndarray, cp.generic))
 
 def is_torch_array(x):
+    """
+    Return True if `x` is a PyTorch tensor.
+
+    This function does not import PyTorch if it has not already been imported
+    and is therefore cheap to use.
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_dask_array
+    is_jax_array
+    """
     # Avoid importing torch if it isn't already
     if 'torch' not in sys.modules:
         return False
@@ -49,6 +101,22 @@ def is_torch_array(x):
     return isinstance(x, torch.Tensor)
 
 def is_dask_array(x):
+    """
+    Return True if `x` is a dask.array Array.
+
+    This function does not import dask if it has not already been imported
+    and is therefore cheap to use.
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_jax_array
+    """
     # Avoid importing dask if it isn't already
     if 'dask.array' not in sys.modules:
         return False
@@ -58,6 +126,23 @@ def is_dask_array(x):
     return isinstance(x, dask.array.Array)
 
 def is_jax_array(x):
+    """
+    Return True if `x` is a JAX array.
+
+    This function does not import JAX if it has not already been imported
+    and is therefore cheap to use.
+
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    """
     # Avoid importing jax if it isn't already
     if 'jax' not in sys.modules:
         return False
@@ -68,7 +153,17 @@ def is_jax_array(x):
 
 def is_array_api_obj(x):
     """
-    Check if x is an array API compatible array object.
+    Return True if `x` is an array API compatible array object.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
     """
     return is_numpy_array(x) \
         or is_cupy_array(x) \
@@ -83,62 +178,128 @@ def _check_api_version(api_version):
     elif api_version is not None and api_version != '2022.12':
         raise ValueError("Only the 2022.12 version of the array API specification is currently supported")
 
-def array_namespace(*xs, api_version=None, _use_compat=True):
+def array_namespace(*xs, api_version=None, use_compat=None):
     """
     Get the array API compatible namespace for the arrays `xs`.
 
-    `xs` should contain one or more arrays.
+    Parameters
+    ----------
+    xs: arrays
+        one or more arrays.
+
+    api_version: str
+        The newest version of the spec that you need support for (currently
+        the compat library wrapped APIs support v2022.12).
 
-    Typical usage is
+    use_compat: bool or None
+        If None (the default), the native namespace will be returned if it is
+        already array API compatible, otherwise a compat wrapper is used. If
+        True, the compat library wrapped library will be returned. If False,
+        the native library namespace is returned.
 
-        def your_function(x, y):
-            xp = array_api_compat.array_namespace(x, y)
-            # Now use xp as the array library namespace
-            return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)
+    Returns
+    -------
+
+    out: namespace
+        The array API compatible namespace corresponding to the arrays in `xs`.
+
+    Raises
+    ------
+    TypeError
+        If `xs` contains arrays from different array libraries or contains a
+        non-array.
+
+
+    Typical usage is to pass the arguments of a function to
+    `array_namespace()` at the top of a function to get the corresponding
+    array API namespace:
+
+    .. code:: python
+
+       def your_function(x, y):
+           xp = array_api_compat.array_namespace(x, y)
+           # Now use xp as the array library namespace
+           return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)
+
+
+    Wrapped array namespaces can also be imported directly. For example,
+    `array_namespace(np.array(...))` will return `array_api_compat.numpy`.
+    This function will also work for any array library not wrapped by
+    array-api-compat if it explicitly defines `__array_namespace__
+    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__array_namespace__.html>`__
+    (the wrapped namespace is always preferred if it exists).
+
+    See Also
+    --------
+
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
 
-    api_version should be the newest version of the spec that you need support
-    for (currently the compat library wrapped APIs only support v2021.12).
     """
+    if use_compat not in [None, True, False]:
+        raise ValueError("use_compat must be None, True, or False")
+
+    _use_compat = use_compat in [None, True]
+
     namespaces = set()
     for x in xs:
         if is_numpy_array(x):
-            _check_api_version(api_version)
-            if _use_compat:
-                from .. import numpy as numpy_namespace
+            from .. import numpy as numpy_namespace
+            import numpy as np
+            if use_compat is True:
+                _check_api_version(api_version)
                 namespaces.add(numpy_namespace)
-            else:
-                import numpy as np
+            elif use_compat is False:
                 namespaces.add(np)
+            else:
+                # numpy 2.0 has __array_namespace__ and is fully array API
+                # compatible.
+                if hasattr(x, '__array_namespace__'):
+                    namespaces.add(x.__array_namespace__(api_version=api_version))
+                else:
+                    namespaces.add(numpy_namespace)
         elif is_cupy_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from .. import cupy as cupy_namespace
                 namespaces.add(cupy_namespace)
             else:
                 import cupy as cp
                 namespaces.add(cp)
         elif is_torch_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from .. import torch as torch_namespace
                 namespaces.add(torch_namespace)
             else:
                 import torch
                 namespaces.add(torch)
         elif is_dask_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from ..dask import array as dask_namespace
                 namespaces.add(dask_namespace)
             else:
-                raise TypeError("_use_compat cannot be False if input array is a dask array!")
+                import dask.array as da
+                namespaces.add(da)
         elif is_jax_array(x):
-            _check_api_version(api_version)
-            # 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
+            if use_compat is True:
+                _check_api_version(api_version)
+                raise ValueError("JAX does not have an array-api-compat wrapper")
+            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
             namespaces.add(jnp)
         elif hasattr(x, '__array_namespace__'):
+            if use_compat is True:
+                raise ValueError("The given array does not have an array-api-compat wrapper")
             namespaces.add(x.__array_namespace__(api_version=api_version))
         else:
             # TODO: Support Python scalars?
@@ -181,15 +342,33 @@ def device(x: Array, /) -> Device:
     """
     Hardware device the array data resides on.
 
+    This is equivalent to `x.device` according to the `standard
+    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.device.html>`__.
+    This helper is included because some array libraries either do not have
+    the `device` attribute or include it with an incompatible API.
+
     Parameters
     ----------
     x: array
-        array instance from NumPy or an array API compatible library.
+        array instance from an array API compatible library.
 
     Returns
     -------
     out: device
-        a ``device`` object (see the "Device Support" section of the array API specification).
+        a ``device`` object (see the `Device Support <https://data-apis.org/array-api/latest/design_topics/device_support.html>`__
+        section of the array API specification).
+
+    Notes
+    -----
+
+    For NumPy the device is always `"cpu"`. For Dask, the device is always a
+    special `DASK_DEVICE` object.
+
+    See Also
+    --------
+
+    to_device : Move array data to a different device.
+
     """
     if is_numpy_array(x):
         return "cpu"
@@ -262,22 +441,50 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
     """
     Copy the array from the device on which it currently resides to the specified ``device``.
 
+    This is equivalent to `x.to_device(device, stream=stream)` according to
+    the `standard
+    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.to_device.html>`__.
+    This helper is included because some array libraries do not have the
+    `to_device` method.
+
     Parameters
     ----------
+
     x: array
-        array instance from NumPy or an array API compatible library.
+        array instance from an array API compatible library.
+
     device: device
-        a ``device`` object (see the "Device Support" section of the array API specification).
+        a ``device`` object (see the `Device Support <https://data-apis.org/array-api/latest/design_topics/device_support.html>`__
+        section of the array API specification).
+
     stream: Optional[Union[int, Any]]
-        stream object to use during copy. In addition to the types supported in ``array.__dlpack__``, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable.
+        stream object to use during copy. In addition to the types supported
+        in ``array.__dlpack__``, implementations may choose to support any
+        library-specific stream object with the caveat that any code using
+        such an object would not be portable.
 
     Returns
     -------
+
     out: array
-        an array with the same data and data type as ``x`` and located on the specified ``device``.
+        an array with the same data and data type as ``x`` and located on the
+        specified ``device``.
+
+    Notes
+    -----
+
+    For NumPy, this function effectively does nothing since the only supported
+    device is the CPU. For CuPy, this method supports CuPy CUDA
+    :external+cupy:class:`Device <cupy.cuda.Device>` and
+    :external+cupy:class:`Stream <cupy.cuda.Stream>` objects. For PyTorch,
+    this is the same as :external+torch:meth:`x.to(device) <torch.Tensor.to>`
+    (the ``stream`` argument is not supported in PyTorch).
+
+    See Also
+    --------
+
+    device : Hardware device the array data resides on.
 
-    .. note::
-       If ``stream`` is given, the copy operation should be enqueued on the provided ``stream``; otherwise, the copy operation should be enqueued on the default stream/queue. Whether the copy is performed synchronously or asynchronously is implementation-dependent. Accordingly, if synchronization is required to guarantee data safety, this must be clearly explained in a conforming library's documentation.
     """
     if is_numpy_array(x):
         if stream is not None:
@@ -305,7 +512,13 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
 
 def size(x):
     """
-    Return the total number of elements of x
+    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>`.
+
     """
     if None in x.shape:
         return None

{array_api_compat-1.5 → array_api_compat-1.6}/array_api_compat/cupy/_aliases.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from functools import partial
-
 import cupy as cp
 
 from ..common import _aliases
 from .._internal import get_xp
 
-asarray = asarray_cupy = partial(_aliases._asarray, namespace='cupy')
-asarray.__doc__ = _aliases._asarray.__doc__
-del partial
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from typing import Optional, Union
+    from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
 
 bool = cp.bool_
 
@@ -62,6 +61,52 @@ matmul = get_xp(cp)(_aliases.matmul)
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
 
+_copy_default = object()
+
+# asarray also adds the copy keyword, which is not present in numpy 1.0.
+def asarray(
+    obj: Union[
+        ndarray,
+        bool,
+        int,
+        float,
+        NestedSequence[bool | int | float],
+        SupportsBufferProtocol,
+    ],
+    /,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    copy: Optional[bool] = _copy_default,
+    **kwargs,
+) -> ndarray:
+    """
+    Array API compatibility wrapper for asarray().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+    with cp.cuda.Device(device):
+        # cupy is like NumPy 1.26 (except without _CopyMode). See the comments
+        # in asarray in numpy/_aliases.py.
+        if copy is not _copy_default:
+            # A future version of CuPy will change the meaning of copy=False
+            # to mean no-copy. We don't know for certain what version it will
+            # be yet, so to avoid breaking that version, we use a different
+            # default value for copy so asarray(obj) with no copy kwarg will
+            # always do the copy-if-needed behavior.
+
+            # This will still need to be updated to remove the
+            # NotImplementedError for copy=False, but at least this won't
+            # break the default or existing behavior.
+            if copy is None:
+                copy = False
+            elif copy is False:
+                raise NotImplementedError("asarray(copy=False) is not yet supported in cupy")
+            kwargs['copy'] = copy
+
+        return cp.array(obj, dtype=dtype, **kwargs)
+
 # 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'):
@@ -73,7 +118,7 @@ if hasattr(cp, 'isdtype'):
 else:
     isdtype = get_xp(cp)(_aliases.isdtype)
 
-__all__ = _aliases.__all__ + ['asarray', 'asarray_cupy', 'bool', 'acos',
+__all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',
                               'bitwise_right_shift', 'concat', 'pow']