Move to Array API version 2023.12.

Author: hameerabbasi

.github/workflows/ci.yml

@@ -154,7 +154,7 @@ jobs:
       uses: actions/checkout@v4
       with:
         repository: data-apis/array-api-tests
-        ref: '33f2d2ea2f3dd2b3ceeeb4519d55e08096184149'  # Latest commit as of 2024-05-29
+        ref: 'd295a0a66cd82a43e84c1b8d73ca198cc45e9d23'  # Latest commit as of 2024-05-29
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python

ci/Numba-array-api-xfails.txt

@@ -29,42 +29,48 @@ array_api_tests/test_has_names.py::test_has_names[linalg-tensordot]
 array_api_tests/test_has_names.py::test_has_names[linalg-trace]
 array_api_tests/test_has_names.py::test_has_names[linalg-vecdot]
 array_api_tests/test_has_names.py::test_has_names[linalg-vector_norm]
+array_api_tests/test_has_names.py::test_has_names[manipulation-repeat]
+array_api_tests/test_has_names.py::test_has_names[manipulation-tile]
 array_api_tests/test_has_names.py::test_has_names[set-unique_all]
 array_api_tests/test_has_names.py::test_has_names[set-unique_inverse]
 array_api_tests/test_has_names.py::test_has_names[creation-arange]
 array_api_tests/test_has_names.py::test_has_names[creation-from_dlpack]
 array_api_tests/test_has_names.py::test_has_names[creation-linspace]
 array_api_tests/test_has_names.py::test_has_names[creation-meshgrid]
+array_api_tests/test_has_names.py::test_has_names[searching-searchsorted]
 array_api_tests/test_has_names.py::test_has_names[sorting-argsort]
+array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
 array_api_tests/test_has_names.py::test_has_names[data_type-isdtype]
 array_api_tests/test_has_names.py::test_has_names[array_method-__dlpack__]
 array_api_tests/test_has_names.py::test_has_names[array_method-__dlpack_device__]
 array_api_tests/test_has_names.py::test_has_names[array_method-__setitem__]
 array_api_tests/test_indexing_functions.py::test_take
 array_api_tests/test_linalg.py::test_vecdot
+array_api_tests/test_manipulation_functions.py::test_repeat
+array_api_tests/test_manipulation_functions.py::test_tile
 array_api_tests/test_operators_and_elementwise_functions.py::test_ceil
 array_api_tests/test_operators_and_elementwise_functions.py::test_trunc
 array_api_tests/test_searching_functions.py::test_argmax
 array_api_tests/test_searching_functions.py::test_argmin
+array_api_tests/test_searching_functions.py::test_searchsorted
 array_api_tests/test_set_functions.py::test_unique_all
 array_api_tests/test_set_functions.py::test_unique_inverse
+array_api_tests/test_statistical_functions.py::test_cumulative_sum
 array_api_tests/test_signatures.py::test_func_signature[unique_all]
 array_api_tests/test_signatures.py::test_func_signature[unique_inverse]
 array_api_tests/test_signatures.py::test_func_signature[arange]
+array_api_tests/test_signatures.py::test_func_signature[cumulative_sum]
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 array_api_tests/test_signatures.py::test_func_signature[linspace]
 array_api_tests/test_signatures.py::test_func_signature[meshgrid]
+array_api_tests/test_signatures.py::test_func_signature[repeat]
+array_api_tests/test_signatures.py::test_func_signature[tile]
 array_api_tests/test_signatures.py::test_func_signature[argsort]
+array_api_tests/test_signatures.py::test_func_signature[searchsorted]
 array_api_tests/test_signatures.py::test_func_signature[isdtype]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack_device__]
 array_api_tests/test_signatures.py::test_array_method_signature[__setitem__]
 array_api_tests/test_sorting_functions.py::test_argsort
 array_api_tests/test_sorting_functions.py::test_sort
-array_api_tests/test_special_cases.py::test_nan_propagation[max]
-array_api_tests/test_special_cases.py::test_nan_propagation[mean]
-array_api_tests/test_special_cases.py::test_nan_propagation[min]
-array_api_tests/test_special_cases.py::test_nan_propagation[prod]
-array_api_tests/test_special_cases.py::test_nan_propagation[std]
-array_api_tests/test_special_cases.py::test_nan_propagation[sum]
-array_api_tests/test_special_cases.py::test_nan_propagation[var]
+array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_sum]

sparse/__init__.py

@@ -4,7 +4,7 @@
 
 from ._version import __version__, __version_tuple__  # noqa: F401
 
-__array_api_version__ = "2022.12"
+__array_api_version__ = "2023.12"
 
 
 class BackendType(Enum):
@@ -45,13 +45,7 @@ def get_backend_module():
 
 
 def __getattr__(attr):
-    if attr == "numba_backend":
-        import sparse.numba_backend as backend_module
-
-        return backend_module
-    if attr == "finch_backend":
-        import sparse.finch_backend as backend_module
-
-        return backend_module
+    if attr == "numba_backend" or attr == "finch_backend":
+        raise AttributeError
 
     return getattr(Backend.get_backend_module(), attr)

sparse/numba_backend/__init__.py

@@ -1,3 +1,5 @@
+import sparse.numba_backend._info as _info
+
 from numpy import (
     add,
     bitwise_and,
@@ -9,6 +11,7 @@
     complex64,
     complex128,
     conj,
+    copysign,
     cos,
     cosh,
     divide,
@@ -23,6 +26,7 @@
     floor_divide,
     greater,
     greater_equal,
+    hypot,
     iinfo,
     inf,
     int8,
@@ -41,6 +45,8 @@
     logical_not,
     logical_or,
     logical_xor,
+    maximum,
+    minimum,
     multiply,
     nan,
     negative,
@@ -50,6 +56,7 @@
     positive,
     remainder,
     sign,
+    signbit,
     sin,
     sinh,
     sqrt,
@@ -119,6 +126,7 @@
     std,
     sum,
     tensordot,
+    unstack,
     var,
     vecdot,
     zeros,
@@ -157,10 +165,16 @@
     where,
 )
 from ._dok import DOK
+from ._info import capabilities, default_device, default_dtypes, devices, dtypes
 from ._io import load_npz, save_npz
 from ._umath import elemwise
 from ._utils import random
 
+
+def __array_namespace_info__():
+    return _info
+
+
 __all__ = [
     "COO",
     "DOK",
@@ -196,19 +210,25 @@
     "broadcast_arrays",
     "broadcast_to",
     "can_cast",
+    "capabilities",
     "ceil",
     "clip",
     "complex128",
     "complex64",
     "concat",
     "concatenate",
     "conj",
+    "copysign",
     "cos",
     "cosh",
+    "default_device",
+    "default_dtypes",
+    "devices",
     "diagonal",
     "diagonalize",
     "divide",
     "dot",
+    "dtypes",
     "e",
     "einsum",
     "elemwise",
@@ -230,6 +250,7 @@
     "full_like",
     "greater",
     "greater_equal",
+    "hypot",
     "iinfo",
     "imag",
     "inf",
@@ -258,8 +279,10 @@
     "matmul",
     "matrix_transpose",
     "max",
+    "maximum",
     "mean",
     "min",
+    "minimum",
     "moveaxis",
     "multiply",
     "nan",
@@ -291,6 +314,7 @@
     "round",
     "save_npz",
     "sign",
+    "signbit",
     "sin",
     "sinh",
     "sort",
@@ -314,6 +338,7 @@
     "uint8",
     "unique_counts",
     "unique_values",
+    "unstack",
     "var",
     "vecdot",
     "where",

sparse/numba_backend/_common.py

@@ -155,7 +155,7 @@ def tensordot(a, b, axes=2, *, return_type=None):
     bs = b.shape
     ndb = b.ndim
     equal = True
-    if nda == 0 or ndb == 0:
+    if not (builtins.all(-nda <= ax < nda for ax in axes_a) and builtins.all(-ndb <= ax < ndb for ax in axes_b)):
         pos = int(nda != 0)
         raise ValueError(f"Input {pos} operand does not have enough dimensions")
     if na != nb:
@@ -2146,10 +2146,22 @@ def reshape(x, /, shape, *, copy=None):
     return x.reshape(shape=shape)
 
 
-def astype(x, dtype, /, *, copy=True):
+@_check_device
+def astype(x, dtype, /, *, copy=True, device=None):
     return x.astype(dtype, copy=copy)
 
 
+def unstack(x, /, *, axis=0):
+    axis = normalize_axis(axis, x.ndim)
+    out = []
+
+    for i in range(x.shape[axis]):
+        idx = (slice(None),) * axis + (i,)
+        out.append(x[idx])
+
+    return tuple(out)
+
+
 @_support_numpy
 def squeeze(x, /, axis=None):
     """Remove singleton dimensions from array.

sparse/numba_backend/_coo/common.py

@@ -1011,7 +1011,7 @@ def _diagonal_idx(coordlist, axis1, axis2, offset):
     return np.array([i for i in range(len(coordlist[axis1])) if coordlist[axis1][i] + offset == coordlist[axis2][i]])
 
 
-def clip(a, a_min=None, a_max=None, out=None):
+def clip(a, min=None, max=None, out=None):
     """
     Clip (limit) the values in the array.
 
@@ -1042,19 +1042,19 @@ def clip(a, a_min=None, a_max=None, out=None):
     --------
     >>> import sparse
     >>> x = sparse.COO.from_numpy([0, 0, 0, 1, 2, 3])
-    >>> sparse.clip(x, a_min=1).todense()  # doctest: +NORMALIZE_WHITESPACE
+    >>> sparse.clip(x, min=1).todense()  # doctest: +NORMALIZE_WHITESPACE
     array([1, 1, 1, 1, 2, 3])
-    >>> sparse.clip(x, a_max=1).todense()  # doctest: +NORMALIZE_WHITESPACE
+    >>> sparse.clip(x, max=1).todense()  # doctest: +NORMALIZE_WHITESPACE
     array([0, 0, 0, 1, 1, 1])
-    >>> sparse.clip(x, a_min=1, a_max=2).todense()  # doctest: +NORMALIZE_WHITESPACE
+    >>> sparse.clip(x, min=1, max=2).todense()  # doctest: +NORMALIZE_WHITESPACE
     array([1, 1, 1, 1, 2, 2])
 
     See Also
     --------
     numpy.clip : Equivalent NumPy function
     """
     a = asCOO(a, name="clip")
-    return a.clip(a_min, a_max)
+    return a.clip(min, max)
 
 
 def expand_dims(x, /, *, axis=0):

sparse/numba_backend/_info.py

@@ -0,0 +1,95 @@
+import numpy as np
+
+from ._common import _check_device
+
+__all__ = [
+    "capabilities",
+    "default_device",
+    "default_dtypes",
+    "devices",
+    "dtypes",
+]
+
+_CAPABILITIES = {
+    "boolean indexing": True,
+    "data-dependent shapes": True,
+}
+
+_DEFAULT_DTYPES = {
+    "cpu": {
+        "real floating": np.dtype(np.float64),
+        "complex floating": np.dtype(np.complex128),
+        "integral": np.dtype(np.int64),
+        "indexing": np.dtype(np.int64),
+    }
+}
+
+
+def _get_dtypes_with_prefix(prefix: str):
+    out = set()
+    for a in np.__all__:
+        if not a.startswith(prefix):
+            continue
+        try:
+            dt = np.dtype(getattr(np, a))
+            out.add(dt)
+        except (ValueError, TypeError, AttributeError):
+            pass
+    return sorted(out)
+
+
+_DTYPES = {
+    "cpu": {
+        "bool": [np.bool_],
+        "signed integer": _get_dtypes_with_prefix("int"),
+        "unsigned integer": _get_dtypes_with_prefix("uint"),
+        "real floating": _get_dtypes_with_prefix("float"),
+        "complex floating": _get_dtypes_with_prefix("complex"),
+    }
+}
+
+for _dtdict in _DTYPES.values():
+    _dtdict["integral"] = _dtdict["signed integer"] + _dtdict["unsigned integer"]
+    _dtdict["numeric"] = _dtdict["integral"] + _dtdict["real floating"] + _dtdict["complex floating"]
+
+del _dtdict
+
+
+def capabilities():
+    return _CAPABILITIES
+
+
+def default_device():
+    return "cpu"
+
+
+@_check_device
+def default_dtypes(*, device=None):
+    if device is None:
+        device = default_device()
+    return _DEFAULT_DTYPES[device]
+
+
+def devices():
+    return ["cpu"]
+
+
+@_check_device
+def dtypes(*, device=None, kind=None):
+    if device is None:
+        device = default_device()
+
+    device_dtypes = _DTYPES[device]
+
+    if kind is None:
+        return device_dtypes
+
+    if isinstance(kind, str):
+        return device_dtypes[kind]
+
+    out = {}
+
+    for k in kind:
+        out[k] = device_dtypes[k]
+
+    return out

sparse/numba_backend/_sparse_array.py

@@ -607,10 +607,12 @@ def clip(self, min=None, max=None, out=None):
         sparse.clip : For full documentation and more details.
         numpy.clip : Equivalent NumPy function.
         """
-        if min is None and max is None:
-            raise ValueError("One of max or min must be given.")
         if out is not None and not isinstance(out, tuple):
             out = (out,)
+        if min is None and max is None:
+            if out is not None:
+                return self.__array_ufunc__(np.identity, "__call__", self, out=out)
+            return self
         return self.__array_ufunc__(np.clip, "__call__", self, a_min=min, a_max=max, out=out)
 
     def astype(self, dtype, casting="unsafe", copy=True):