diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 35a74dd4..dc5ef256 100644
--- a/.github/workflows/ci.yml
+++ b/.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
diff --git a/.gitignore b/.gitignore
index 5621fa44..ec62754d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -80,3 +80,6 @@ sandbox.py
# Version file
sparse/_version.py
+
+# Benchmarks
+/results/
diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt
index b7d1473c..f699bdf2 100644
--- a/ci/Numba-array-api-xfails.txt
+++ b/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]
diff --git a/sparse/__init__.py b/sparse/__init__.py
index 6a332cf2..ef16a508 100644
--- a/sparse/__init__.py
+++ b/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"
diff --git a/sparse/numba_backend/__init__.py b/sparse/numba_backend/__init__.py
index 4658109b..ff0d9968 100644
--- a/sparse/numba_backend/__init__.py
+++ b/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",
diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py
index ac396904..b5ec86e7 100644
--- a/sparse/numba_backend/_common.py
+++ b/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.
diff --git a/sparse/numba_backend/_coo/common.py b/sparse/numba_backend/_coo/common.py
index dc831a6d..c35e98f0 100644
--- a/sparse/numba_backend/_coo/common.py
+++ b/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,11 +1042,11 @@ 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
@@ -1054,7 +1054,7 @@ def clip(a, a_min=None, a_max=None, out=None):
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):
diff --git a/sparse/numba_backend/_info.py b/sparse/numba_backend/_info.py
new file mode 100644
index 00000000..6d66e2c0
--- /dev/null
+++ b/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
diff --git a/sparse/numba_backend/_sparse_array.py b/sparse/numba_backend/_sparse_array.py
index 763a779e..aa0512d5 100644
--- a/sparse/numba_backend/_sparse_array.py
+++ b/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):
diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py
index ed1ef96f..b107ec26 100644
--- a/sparse/numba_backend/tests/test_coo.py
+++ b/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
diff --git a/sparse/numba_backend/tests/test_namespace.py b/sparse/numba_backend/tests/test_namespace.py
index 39556f99..06f61b39 100644
--- a/sparse/numba_backend/tests/test_namespace.py
+++ b/sparse/numba_backend/tests/test_namespace.py
@@ -36,6 +36,7 @@ def test_namespace():
"bool",
"broadcast_arrays",
"broadcast_to",
+ "capabilities",
"can_cast",
"ceil",
"clip",
@@ -44,8 +45,13 @@ def test_namespace():
"concat",
"concatenate",
"conj",
+ "copysign",
"cos",
"cosh",
+ "default_device",
+ "default_dtypes",
+ "devices",
+ "dtypes",
"diagonal",
"diagonalize",
"divide",
@@ -71,6 +77,7 @@ def test_namespace():
"full_like",
"greater",
"greater_equal",
+ "hypot",
"iinfo",
"imag",
"inf",
@@ -99,8 +106,10 @@ def test_namespace():
"matrix_transpose",
"matmul",
"max",
+ "maximum",
"mean",
"min",
+ "minimum",
"moveaxis",
"multiply",
"nan",
@@ -132,6 +141,7 @@ def test_namespace():
"round",
"save_npz",
"sign",
+ "signbit",
"sin",
"sinh",
"sort",
@@ -155,6 +165,7 @@ def test_namespace():
"uint8",
"unique_counts",
"unique_values",
+ "unstack",
"var",
"vecdot",
"where",