Move to Array API version 2023.12.

.github/workflows/ci.yml

@@ -155,7 +155,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-06-10
  submodules: 'true'
  path: 'array-api-tests'
  - name: Set up Python

.gitignore

@@ -80,3 +80,6 @@ sandbox.py
 
 # Version file
 sparse/_version.py
+
+# Benchmarks
+/results/

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,8 +4,6 @@
 
 from ._version import __version__, __version_tuple__ # noqa: F401
 
-__array_api_version__ = "2022.12"
-
 
 class BackendType(Enum):
  Numba = "Numba"

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,
@@ -161,7 +169,16 @@
 from ._umath import elemwise
 from ._utils import random
 
+
+def __array_namespace_info__():
+ return _info
+
+
+__array_api_version__ = "2023.12"
+
+
 __all__ = [
+ "__array_api_version__",
  "COO",
  "DOK",
  "GCXS",
@@ -203,6 +220,7 @@
  "concat",
  "concatenate",
  "conj",
+ "copysign",
  "cos",
  "cosh",
  "diagonal",
@@ -230,6 +248,7 @@
  "full_like",
  "greater",
  "greater_equal",
+ "hypot",
  "iinfo",
  "imag",
  "inf",
@@ -258,8 +277,10 @@
  "matmul",
  "matrix_transpose",
  "max",
+ "maximum",
  "mean",
  "min",
+ "minimum",
  "moveaxis",
  "multiply",
  "nan",
@@ -291,6 +312,7 @@
  "round",
  "save_npz",
  "sign",
+ "signbit",
  "sin",
  "sinh",
  "sort",
@@ -314,6 +336,7 @@
  "uint8",
  "unique_counts",
  "unique_values",
+ "unstack",
  "var",
  "vecdot",
  "where",

sparse/numba_backend/_common.py

@@ -155,7 +155,7 @@
  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 @@
  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 @@
  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):

sparse/numba_backend/tests/test_coo.py

@@ -1200,9 +1200,6 @@ def test_clip():
 
  assert_eq(np.clip(s, 1, 3), np.clip(x, 1, 3))
 
- with pytest.raises(ValueError):
- s.clip()
-
  out = sparse.COO.from_numpy(np.zeros_like(x))
  out2 = s.clip(min=1, max=3, out=out)
  assert out is out2