Merge branch 'main' into signbit-nan

Author: asmeurer

.github/workflows/array-api-tests-numpy-1-26.yml

@@ -0,0 +1,11 @@
+name: Array API Tests (NumPy 1.26)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-numpy-latest:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: numpy
+      package-version: '== 1.26.*'
+      xfails-file-extra: '-1-26'

.github/workflows/array-api-tests-numpy.yml renamed to .github/workflows/array-api-tests-numpy-latest.yml

@@ -13,7 +13,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Download Artifact
-        uses: dawidd6/action-download-artifact@v3
+        uses: dawidd6/action-download-artifact@v6
         with:
           workflow: docs-build.yml
           name: docs-build

.github/workflows/docs-deploy.yml

@@ -94,7 +94,7 @@ jobs:
       if: >-
         (github.event_name == 'push' && startsWith(github.ref, 'refs/tags'))
         || (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true')
-      uses: pypa/gh-action-pypi-publish@v1.8.14
+      uses: pypa/gh-action-pypi-publish@v1.9.0
       with:
         repository-url: https://test.pypi.org/legacy/
         print-hash: true
@@ -107,6 +107,6 @@ jobs:
 
     - name: Publish distribution 📦 to PyPI
       if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
-      uses: pypa/gh-action-pypi-publish@v1.8.14
+      uses: pypa/gh-action-pypi-publish@v1.9.0
       with:
         print-hash: true

.github/workflows/publish-package.yml

@@ -6,7 +6,7 @@ jobs:
     strategy:
       matrix:
         python-version: ['3.9', '3.10', '3.11', '3.12']
-        numpy-version: ['1.21', '1.26', 'dev']
+        numpy-version: ['1.21', '1.26', '2.0', 'dev']
         exclude:
           - python-version: '3.11'
             numpy-version: '1.21'
@@ -28,6 +28,11 @@ jobs:
           else
             PIP_EXTRA='numpy==1.26.*'
           fi
+
+          if [ "${{ matrix.python-version }}" == "3.9" ]; then
+            sed -i '/^ndonnx/d' requirements-dev.txt
+          fi
+
           python -m pip install -r requirements-dev.txt $PIP_EXTRA
 
       - name: Run Tests

.github/workflows/tests.yml

@@ -2,8 +2,8 @@
 
 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, JAX and `sparse` 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).
+NumPy, CuPy, PyTorch, Dask, JAX, ndonnx and `sparse` 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/

README.md

@@ -17,6 +17,6 @@
 this implementation for the default when working with NumPy arrays.
 
 """
-__version__ = '1.6'
+__version__ = '1.8'
 
 from .common import *  # noqa: F401, F403

array_api_compat/__init__.py

@@ -12,7 +12,7 @@
 from typing import NamedTuple
 import inspect
 
-from ._helpers import _check_device
+from ._helpers import array_namespace, _check_device
 
 # These functions are modified from the NumPy versions.
 
@@ -264,6 +264,66 @@ def var(
 ) -> ndarray:
     return xp.var(x, axis=axis, ddof=correction, keepdims=keepdims, **kwargs)
 
+
+# The min and max argument names in clip are different and not optional in numpy, and type
+# promotion behavior is different.
+def clip(
+    x: ndarray,
+    /,
+    min: Optional[Union[int, float, ndarray]] = None,
+    max: Optional[Union[int, float, ndarray]] = None,
+    *,
+    xp,
+    # TODO: np.clip has other ufunc kwargs
+    out: Optional[ndarray] = None,
+) -> ndarray:
+    def _isscalar(a):
+        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)
+
+    # 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.,
+    # avoiding uint64 -> float64 promotion).
+
+    # Note: cases where min or max overflow (integer) or round (float) in the
+    # wrong direction when downcasting to x.dtype are unspecified. This code
+    # just does whatever NumPy does when it downcasts in the assignment, but
+    # other behavior could be preferred, especially for integers. For example,
+    # this code produces:
+
+    # >>> clip(asarray(0, dtype=int8), asarray(128, dtype=int16), None)
+    # -128
+
+    # but an answer of 0 might be preferred. See
+    # https://github.com/numpy/numpy/issues/24976 for more discussion on this issue.
+
+
+    # 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:
+        min = None
+    if type(max) is int and max >= xp.iinfo(x.dtype).max:
+        max = None
+
+    if out is None:
+        out = wrapped_xp.asarray(xp.broadcast_to(x, result_shape), copy=True)
+    if min is not None:
+        a = xp.broadcast_to(xp.asarray(min), 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)
+        ib = (out > b) | xp.isnan(b)
+        out[ib] = wrapped_xp.astype(b[ib], out.dtype)
+    # Return a scalar for 0-D
+    return out[()]
+
 # Unlike transpose(), the axes argument to permute_dims() is required.
 def permute_dims(x: ndarray, /, axes: Tuple[int, ...], xp) -> ndarray:
     return xp.transpose(x, axes)
@@ -465,6 +525,6 @@ def isdtype(
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
-           'astype', 'std', 'var', 'permute_dims', 'reshape', 'argsort',
+           'astype', 'std', 'var', 'clip', 'permute_dims', 'reshape', 'argsort',
            'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
            'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/common/_aliases.py

@@ -48,9 +48,10 @@ def is_numpy_array(x):
     is_array_api_obj
     is_cupy_array
     is_torch_array
+    is_ndonnx_array
     is_dask_array
     is_jax_array
-    is_pydata_sparse
+    is_pydata_sparse_array
     """
     # Avoid importing NumPy if it isn't already
     if 'numpy' not in sys.modules:
@@ -78,11 +79,12 @@ def is_cupy_array(x):
     is_array_api_obj
     is_numpy_array
     is_torch_array
+    is_ndonnx_array
     is_dask_array
     is_jax_array
-    is_pydata_sparse
+    is_pydata_sparse_array
     """
-    # Avoid importing NumPy if it isn't already
+    # Avoid importing CuPy if it isn't already
     if 'cupy' not in sys.modules:
         return False
 
@@ -107,7 +109,7 @@ def is_torch_array(x):
     is_cupy_array
     is_dask_array
     is_jax_array
-    is_pydata_sparse
+    is_pydata_sparse_array
     """
     # Avoid importing torch if it isn't already
     if 'torch' not in sys.modules:
@@ -118,6 +120,33 @@ def is_torch_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, torch.Tensor)
 
+def is_ndonnx_array(x):
+    """
+    Return True if `x` is a ndonnx Array.
+
+    This function does not import ndonnx 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_ndonnx_array
+    is_dask_array
+    is_jax_array
+    is_pydata_sparse_array
+    """
+    # Avoid importing torch if it isn't already
+    if 'ndonnx' not in sys.modules:
+        return False
+
+    import ndonnx as ndx
+
+    return isinstance(x, ndx.Array)
+
 def is_dask_array(x):
     """
     Return True if `x` is a dask.array Array.
@@ -133,8 +162,9 @@ def is_dask_array(x):
     is_numpy_array
     is_cupy_array
     is_torch_array
+    is_ndonnx_array
     is_jax_array
-    is_pydata_sparse
+    is_pydata_sparse_array
     """
     # Avoid importing dask if it isn't already
     if 'dask.array' not in sys.modules:
@@ -160,8 +190,9 @@ def is_jax_array(x):
     is_numpy_array
     is_cupy_array
     is_torch_array
+    is_ndonnx_array
     is_dask_array
-    is_pydata_sparse
+    is_pydata_sparse_array
     """
     # Avoid importing jax if it isn't already
     if 'jax' not in sys.modules:
@@ -172,7 +203,7 @@ def is_jax_array(x):
     return isinstance(x, jax.Array) or _is_jax_zero_gradient_array(x)
 
 
-def is_pydata_sparse(x) -> bool:
+def is_pydata_sparse_array(x) -> bool:
     """
     Return True if `x` is an array from the `sparse` package.
 
@@ -188,6 +219,7 @@ def is_pydata_sparse(x) -> bool:
     is_numpy_array
     is_cupy_array
     is_torch_array
+    is_ndonnx_array
     is_dask_array
     is_jax_array
     """
@@ -211,6 +243,7 @@ def is_array_api_obj(x):
     is_numpy_array
     is_cupy_array
     is_torch_array
+    is_ndonnx_array
     is_dask_array
     is_jax_array
     """
@@ -219,7 +252,7 @@ def is_array_api_obj(x):
         or is_torch_array(x) \
         or is_dask_array(x) \
         or is_jax_array(x) \
-        or is_pydata_sparse(x) \
+        or is_pydata_sparse_array(x) \
         or hasattr(x, '__array_namespace__')
 
 def _check_api_version(api_version):
@@ -288,7 +321,7 @@ def your_function(x, y):
     is_torch_array
     is_dask_array
     is_jax_array
-    is_pydata_sparse
+    is_pydata_sparse_array
 
     """
     if use_compat not in [None, True, False]:
@@ -307,12 +340,9 @@ def your_function(x, y):
             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)
@@ -344,11 +374,15 @@ def your_function(x, y):
             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(x):
+        elif is_pydata_sparse_array(x):
             if use_compat is True:
                 _check_api_version(api_version)
                 raise ValueError("`sparse` does not have an array-api-compat wrapper")
@@ -451,7 +485,7 @@ def device(x: Array, /) -> Device:
             return x.device()
         else:
             return x.device
-    elif is_pydata_sparse(x):
+    elif is_pydata_sparse_array(x):
         # `sparse` will gain `.device`, so check for this first.
         x_device = getattr(x, 'device', None)
         if x_device is not None:
@@ -580,10 +614,11 @@ 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(x) and device == _device(x):
+    elif is_pydata_sparse_array(x) and device == _device(x):
         # Perform trivial check to return the same array if
         # device is same instead of err-ing.
         return x
@@ -613,7 +648,8 @@ def size(x):
     "is_jax_array",
     "is_numpy_array",
     "is_torch_array",
-    "is_pydata_sparse",
+    "is_ndonnx_array",
+    "is_pydata_sparse_array",
     "size",
     "to_device",
 ]

array_api_compat/common/_helpers.py

@@ -47,6 +47,7 @@
 astype = _aliases.astype
 std = get_xp(cp)(_aliases.std)
 var = get_xp(cp)(_aliases.var)
+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)