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",