Repository: pydata/sparse
Pull request #691 (2024-05-30 06:30:03)

Expand Array API coverage for Numba backend

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pydata/sparse/pull/691?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pydata) Report
Attention: Patch coverage is `89.65517%` with `3 lines` in your changes are missing coverage. Please review.
> Project coverage is 90.02%. Comparing base [(`cb6b604`)](https://app.codecov.io/gh/pydata/sparse/commit/cb6b604698f477cc677bce302370ab530cdd0b4b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pydata) to head [(`a507b60`)](https://app.codecov.io/gh/pydata/sparse/commit/a507b60c8dbbf28efc57296eb3564e699a78782a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pydata).
> Report is 56 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #691      +/-   ##
==========================================
- Coverage   90.21%   90.02%   -0.19%     
==========================================
  Files          20       22       +2     
  Lines        3670     3721      +51     
==========================================
+ Hits         3311     3350      +39     
- Misses        359      371      +12     
```



</details>

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 025026d..0a88c6d 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -173,7 +173,7 @@ jobs:
         SPARSE_BACKEND: ${{ matrix.backend }}
       run: |
         cd ${GITHUB_WORKSPACE}/array-api-tests
-        pytest array_api_tests -v -c pytest.ini -n 4 --ci --max-examples=2 --derandomize --disable-deadline -o xfail_strict=True --xfails-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-xfails.txt --skips-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-skips.txt
+        pytest array_api_tests -v -c pytest.ini -n 4 --max-examples=2 --derandomize --disable-deadline -o xfail_strict=True --xfails-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-xfails.txt --skips-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-skips.txt
 on:
   # Trigger the workflow on push or pull request,
   # but only for the main branch
diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt
index c3de197..b7d1473 100644
--- a/ci/Numba-array-api-xfails.txt
+++ b/ci/Numba-array-api-xfails.txt
@@ -37,68 +37,30 @@ 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[sorting-argsort]
 array_api_tests/test_has_names.py::test_has_names[data_type-isdtype]
-array_api_tests/test_has_names.py::test_has_names[elementwise-conj]
 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_has_names.py::test_has_names[array_method-to_device]
-array_api_tests/test_has_names.py::test_has_names[array_attribute-device]
-array_api_tests/test_has_names.py::test_has_names[array_attribute-mT]
 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_concat
-array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_add[__iadd__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x, s)]
 array_api_tests/test_operators_and_elementwise_functions.py::test_ceil
-array_api_tests/test_operators_and_elementwise_functions.py::test_conj
-array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__truediv__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_imag
-array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__mul__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__imul__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__pow__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__ipow__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_real
-array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub__(x, s)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x, s)]
 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_nonzero_zerodim_error
 array_api_tests/test_set_functions.py::test_unique_all
 array_api_tests/test_set_functions.py::test_unique_inverse
 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[empty]
-array_api_tests/test_signatures.py::test_func_signature[empty_like]
-array_api_tests/test_signatures.py::test_func_signature[eye]
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
-array_api_tests/test_signatures.py::test_func_signature[full]
-array_api_tests/test_signatures.py::test_func_signature[full_like]
 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[ones]
-array_api_tests/test_signatures.py::test_func_signature[ones_like]
-array_api_tests/test_signatures.py::test_func_signature[zeros]
-array_api_tests/test_signatures.py::test_func_signature[zeros_like]
-array_api_tests/test_signatures.py::test_func_signature[broadcast_to]
-array_api_tests/test_signatures.py::test_func_signature[squeeze]
 array_api_tests/test_signatures.py::test_func_signature[argsort]
-array_api_tests/test_signatures.py::test_func_signature[sort]
 array_api_tests/test_signatures.py::test_func_signature[isdtype]
-array_api_tests/test_signatures.py::test_func_signature[conj]
 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_signatures.py::test_array_method_signature[to_device]
 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_unary[isfinite(x_i is NaN) -> False]
-array_api_tests/test_special_cases.py::test_empty_arrays[prod]
 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]
@@ -106,4 +68,3 @@ 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_statistical_functions.py::test_mean
diff --git a/docs/generated/sparse.expand_dims.rst b/docs/generated/sparse.expand_dims.rst
new file mode 100644
index 0000000..7f91864
--- /dev/null
+++ b/docs/generated/sparse.expand_dims.rst
@@ -0,0 +1,6 @@
+expand\_dims
+============
+
+.. currentmodule:: sparse
+
+.. autofunction:: expand_dims
diff --git a/docs/generated/sparse.flip.rst b/docs/generated/sparse.flip.rst
new file mode 100644
index 0000000..9c83383
--- /dev/null
+++ b/docs/generated/sparse.flip.rst
@@ -0,0 +1,6 @@
+flip
+====
+
+.. currentmodule:: sparse
+
+.. autofunction:: flip
diff --git a/docs/generated/sparse.isfinite.rst b/docs/generated/sparse.isfinite.rst
deleted file mode 100644
index b5f7bb9..0000000
--- a/docs/generated/sparse.isfinite.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-isfinite
-========
-
-.. currentmodule:: sparse
-
-.. autofunction:: isfinite
diff --git a/docs/generated/sparse.matrix_transpose.rst b/docs/generated/sparse.matrix_transpose.rst
new file mode 100644
index 0000000..755521d
--- /dev/null
+++ b/docs/generated/sparse.matrix_transpose.rst
@@ -0,0 +1,6 @@
+matrix\_transpose
+=================
+
+.. currentmodule:: sparse
+
+.. autofunction:: matrix_transpose
diff --git a/docs/generated/sparse.rst b/docs/generated/sparse.rst
index 8c8e9b3..9592cc2 100644
--- a/docs/generated/sparse.rst
+++ b/docs/generated/sparse.rst
@@ -80,8 +80,6 @@ API
 
     full_like
 
-    isfinite
-
     isinf
 
     isnan
diff --git a/docs/generated/sparse.sort.rst b/docs/generated/sparse.sort.rst
new file mode 100644
index 0000000..cf1e71a
--- /dev/null
+++ b/docs/generated/sparse.sort.rst
@@ -0,0 +1,6 @@
+sort
+====
+
+.. currentmodule:: sparse
+
+.. autofunction:: sort
diff --git a/docs/generated/sparse.take.rst b/docs/generated/sparse.take.rst
new file mode 100644
index 0000000..42f5ac8
--- /dev/null
+++ b/docs/generated/sparse.take.rst
@@ -0,0 +1,6 @@
+take
+====
+
+.. currentmodule:: sparse
+
+.. autofunction:: take
diff --git a/docs/generated/sparse.var.rst b/docs/generated/sparse.var.rst
new file mode 100644
index 0000000..f8badb5
--- /dev/null
+++ b/docs/generated/sparse.var.rst
@@ -0,0 +1,6 @@
+var
+===
+
+.. currentmodule:: sparse
+
+.. autofunction:: var
diff --git a/docs/generated/sparse.vecdot.rst b/docs/generated/sparse.vecdot.rst
new file mode 100644
index 0000000..995d810
--- /dev/null
+++ b/docs/generated/sparse.vecdot.rst
@@ -0,0 +1,6 @@
+vecdot
+======
+
+.. currentmodule:: sparse
+
+.. autofunction:: vecdot
diff --git a/sparse/numba_backend/__init__.py b/sparse/numba_backend/__init__.py
index b51cbfe..4658109 100644
--- a/sparse/numba_backend/__init__.py
+++ b/sparse/numba_backend/__init__.py
@@ -8,6 +8,7 @@ from numpy import (
     ceil,
     complex64,
     complex128,
+    conj,
     cos,
     cosh,
     divide,
@@ -23,12 +24,12 @@ from numpy import (
     greater,
     greater_equal,
     iinfo,
-    imag,
     inf,
     int8,
     int16,
     int32,
     int64,
+    isfinite,
     less,
     less_equal,
     log,
@@ -47,7 +48,6 @@ from numpy import (
     not_equal,
     pi,
     positive,
-    real,
     remainder,
     sign,
     sin,
@@ -96,7 +96,7 @@ from ._common import (
     eye,
     full,
     full_like,
-    isfinite,
+    imag,
     isinf,
     isnan,
     matmul,
@@ -111,6 +111,7 @@ from ._common import (
     pad,
     permute_dims,
     prod,
+    real,
     reshape,
     round,
     squeeze,
@@ -201,6 +202,7 @@ __all__ = [
     "complex64",
     "concat",
     "concatenate",
+    "conj",
     "cos",
     "cosh",
     "diagonal",
diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py
index 8364467..ac39690 100644
--- a/sparse/numba_backend/_common.py
+++ b/sparse/numba_backend/_common.py
@@ -29,6 +29,17 @@ def _is_scipy_sparse_obj(x):
     return False
 
 
+def _check_device(func):
+    @wraps(func)
+    def wrapped(*args, **kwargs):
+        device = kwargs.get("device", None)
+        if device not in {"cpu", None}:
+            raise ValueError("Device must be `'cpu'` or `None`.")
+        return func(*args, **kwargs)
+
+    return wrapped
+
+
 def _is_sparse(x):
     """
     Tests if the supplied argument is a SciPy sparse object, or one from this library.
@@ -1533,7 +1544,8 @@ def concatenate(arrays, axis=0, compressed_axes=None):
 concat = concatenate
 
 
-def eye(N, M=None, k=0, dtype=float, format="coo", **kwargs):
+@_check_device
+def eye(N, M=None, k=0, dtype=float, format="coo", *, device=None, **kwargs):
     """Return a 2-D array in the specified format with ones on the diagonal and zeros elsewhere.
 
     Parameters
@@ -1595,7 +1607,8 @@ def eye(N, M=None, k=0, dtype=float, format="coo", **kwargs):
     return COO(coords, data=data, shape=(N, M), has_duplicates=False, sorted=True).asformat(format, **kwargs)
 
 
-def full(shape, fill_value, dtype=None, format="coo", order="C", **kwargs):
+@_check_device
+def full(shape, fill_value, dtype=None, format="coo", order="C", *, device=None, **kwargs):
     """Return a SparseArray of given shape and type, filled with `fill_value`.
 
     Parameters
@@ -1649,7 +1662,8 @@ def full(shape, fill_value, dtype=None, format="coo", order="C", **kwargs):
     ).asformat(format, **kwargs)
 
 
-def full_like(a, fill_value, dtype=None, shape=None, format=None, **kwargs):
+@_check_device
+def full_like(a, fill_value, dtype=None, shape=None, format=None, *, device=None, **kwargs):
     """Return a full array with the same shape and type as a given array.
 
     Parameters
@@ -1692,7 +1706,7 @@ def full_like(a, fill_value, dtype=None, shape=None, format=None, **kwargs):
     )
 
 
-def zeros(shape, dtype=float, format="coo", **kwargs):
+def zeros(shape, dtype=float, format="coo", *, device=None, **kwargs):
     """Return a SparseArray of given shape and type, filled with zeros.
 
     Parameters
@@ -1721,10 +1735,10 @@ def zeros(shape, dtype=float, format="coo", **kwargs):
     array([[0, 0],
            [0, 0]])
     """
-    return full(shape, fill_value=0, dtype=np.dtype(dtype), format=format, **kwargs)
+    return full(shape, fill_value=0, dtype=np.dtype(dtype), format=format, device=device, **kwargs)
 
 
-def zeros_like(a, dtype=None, shape=None, format=None, **kwargs):
+def zeros_like(a, dtype=None, shape=None, format=None, *, device=None, **kwargs):
     """Return a SparseArray of zeros with the same shape and type as ``a``.
 
     Parameters
@@ -1750,10 +1764,10 @@ def zeros_like(a, dtype=None, shape=None, format=None, **kwargs):
     array([[0, 0, 0],
            [0, 0, 0]])
     """
-    return full_like(a, fill_value=0, dtype=dtype, shape=shape, format=format, **kwargs)
+    return full_like(a, fill_value=0, dtype=dtype, shape=shape, format=format, device=device, **kwargs)
 
 
-def ones(shape, dtype=float, format="coo", **kwargs):
+def ones(shape, dtype=float, format="coo", *, device=None, **kwargs):
     """Return a SparseArray of given shape and type, filled with ones.
 
     Parameters
@@ -1782,10 +1796,10 @@ def ones(shape, dtype=float, format="coo", **kwargs):
     array([[1, 1],
            [1, 1]])
     """
-    return full(shape, fill_value=1, dtype=np.dtype(dtype), format=format, **kwargs)
+    return full(shape, fill_value=1, dtype=np.dtype(dtype), format=format, device=device, **kwargs)
 
 
-def ones_like(a, dtype=None, shape=None, format=None, **kwargs):
+def ones_like(a, dtype=None, shape=None, format=None, *, device=None, **kwargs):
     """Return a SparseArray of ones with the same shape and type as ``a``.
 
     Parameters
@@ -1811,18 +1825,18 @@ def ones_like(a, dtype=None, shape=None, format=None, **kwargs):
     array([[1, 1, 1],
            [1, 1, 1]])
     """
-    return full_like(a, fill_value=1, dtype=dtype, shape=shape, format=format, **kwargs)
+    return full_like(a, fill_value=1, dtype=dtype, shape=shape, format=format, device=device, **kwargs)
 
 
-def empty(shape, dtype=float, format="coo", **kwargs):
-    return full(shape, fill_value=0, dtype=np.dtype(dtype), format=format, **kwargs)
+def empty(shape, dtype=float, format="coo", *, device=None, **kwargs):
+    return full(shape, fill_value=0, dtype=np.dtype(dtype), format=format, device=device, **kwargs)
 
 
 empty.__doc__ = zeros.__doc__
 
 
-def empty_like(a, dtype=None, shape=None, format=None, **kwargs):
-    return full_like(a, fill_value=0, dtype=dtype, shape=shape, format=format, **kwargs)
+def empty_like(a, dtype=None, shape=None, format=None, *, device=None, **kwargs):
+    return full_like(a, fill_value=0, dtype=dtype, shape=shape, format=format, device=device, **kwargs)
 
 
 empty_like.__doc__ = zeros_like.__doc__
@@ -2004,7 +2018,8 @@ def format_to_string(format):
     raise ValueError(f"invalid format: {format}")
 
 
-def asarray(obj, /, *, dtype=None, format="coo", device=None, copy=False):
+@_check_device
+def asarray(obj, /, *, dtype=None, format="coo", copy=False, device=None):
     """
     Convert the input to a sparse array.
 
@@ -2066,6 +2081,7 @@ def _support_numpy(func):
     we want to flag it and dispatch to NumPy.
     """
 
+    @wraps(func)
     def wrapper_func(*args, **kwargs):
         x = args[0]
         if isinstance(x, np.ndarray | np.number):
@@ -2136,6 +2152,20 @@ def astype(x, dtype, /, *, copy=True):
 
 @_support_numpy
 def squeeze(x, /, axis=None):
+    """Remove singleton dimensions from array.
+
+    Parameters
+    ----------
+    x : SparseArray
+        Input array.
+    axis : int or tuple[int, ...], optional
+        The singleton axes to remove. By default all singleton axes are removed.
+
+    Returns
+    -------
+    output : SparseArray
+        Array with singleton dimensions removed.
+    """
     return x.squeeze(axis=axis)
 
 
@@ -2168,14 +2198,18 @@ def isnan(x, /):
     return x.isnan()
 
 
-def isfinite(x, /):
-    return ~isinf(x)
-
-
 def nonzero(x, /):
     return x.nonzero()
 
 
+def imag(x, /):
+    return x.imag
+
+
+def real(x, /):
+    return x.real
+
+
 def vecdot(x1, x2, /, *, axis=-1):
     """
     Computes the (vector) dot product of two arrays.
@@ -2192,6 +2226,10 @@ def vecdot(x1, x2, /, *, axis=-1):
     out : Union[SparseArray, numpy.ndarray]
         Sparse or 0-D array containing dot product.
     """
+    ndmin = builtins.min((x1.ndim, x2.ndim))
+    if not (-ndmin <= axis < ndmin) or x1.shape[axis] != x2.shape[axis]:
+        raise ValueError("Shapes must match along `axis`.")
+
     if np.issubdtype(x1.dtype, np.complexfloating):
         x1 = np.conjugate(x1)
 
diff --git a/sparse/numba_backend/_compressed/compressed.py b/sparse/numba_backend/_compressed/compressed.py
index f92c368..314ca9f 100644
--- a/sparse/numba_backend/_compressed/compressed.py
+++ b/sparse/numba_backend/_compressed/compressed.py
@@ -324,6 +324,16 @@ class GCXS(SparseArray, NDArrayOperatorsMixin):
     def T(self):
         return self.transpose()
 
+    @property
+    def mT(self):
+        if self.ndim < 2:
+            raise ValueError("Cannot compute matrix transpose if `ndim < 2`.")
+
+        axis = list(range(self.ndim))
+        axis[-1], axis[-2] = axis[-2], axis[-1]
+
+        return self.transpose(axis)
+
     def __str__(self):
         summary = (
             f"<GCXS: shape={self.shape}, dtype={self.dtype}, nnz={self.nnz}, fill_value={self.fill_value}, "
diff --git a/sparse/numba_backend/_coo/common.py b/sparse/numba_backend/_coo/common.py
index 5fe8c2c..55e63cd 100644
--- a/sparse/numba_backend/_coo/common.py
+++ b/sparse/numba_backend/_coo/common.py
@@ -160,6 +160,7 @@ def concatenate(arrays, axis=0):
     check_consistent_fill_value(arrays)
 
     if axis is None:
+        axis = 0
         arrays = [x.flatten() for x in arrays]
 
     arrays = [x if isinstance(x, COO) else COO(x) for x in arrays]
@@ -1243,7 +1244,7 @@ def unique_values(x, /):
     return values
 
 
-def sort(x, /, *, axis=-1, descending=False):
+def sort(x, /, *, axis=-1, descending=False, stable=False):
     """
     Returns a sorted copy of an input array ``x``.
 
@@ -1258,6 +1259,8 @@ def sort(x, /, *, axis=-1, descending=False):
         Sort order. If ``True``, the array must be sorted in descending order (by value).
         If ``False``, the array must be sorted in ascending order (by value).
         Default: ``False``.
+    stable : bool
+        Whether the sort is stable. Only ``False`` is supported currently.
 
     Returns
     -------
@@ -1279,12 +1282,14 @@ def sort(x, /, *, axis=-1, descending=False):
     array([ 2, 2, 1, 0, 0, -3])
 
     """
-
     from .._common import moveaxis
     from .core import COO
 
     x = _validate_coo_input(x)
 
+    if stable:
+        raise ValueError("`stable=True` isn't currently supported.")
+
     original_ndim = x.ndim
     if x.ndim == 1:
         x = x[None, :]
@@ -1357,11 +1362,7 @@ def _validate_coo_input(x: Any):
 
 @numba.jit(nopython=True, nogil=True)
 def _sort_coo(
-    coords: np.ndarray,
-    data: np.ndarray,
-    fill_value: float,
-    sort_axis_len: int,
-    descending: bool,
+    coords: np.ndarray, data: np.ndarray, fill_value: float, sort_axis_len: int, descending: bool
 ) -> tuple[np.ndarray, np.ndarray]:
     assert coords.shape[0] == 2
     group_coords = coords[0, :]
diff --git a/sparse/numba_backend/_coo/core.py b/sparse/numba_backend/_coo/core.py
index ae5d039..575cb48 100644
--- a/sparse/numba_backend/_coo/core.py
+++ b/sparse/numba_backend/_coo/core.py
@@ -687,7 +687,7 @@ class COO(SparseArray, NDArrayOperatorsMixin):  # lgtm [py/missing-equals]
     __repr__ = __str__
 
     def _reduce_calc(self, method, axis, keepdims=False, **kwargs):
-        if axis[0] is None:
+        if axis == (None,):
             axis = tuple(range(self.ndim))
         axis = tuple(a if a >= 0 else a + self.ndim for a in axis)
         neg_axis = tuple(ax for ax in range(self.ndim) if ax not in set(axis))
@@ -849,6 +849,16 @@ class COO(SparseArray, NDArrayOperatorsMixin):  # lgtm [py/missing-equals]
         """
         return self.transpose(tuple(range(self.ndim))[::-1])
 
+    @property
+    def mT(self):
+        if self.ndim < 2:
+            raise ValueError("Cannot compute matrix transpose if `ndim < 2`.")
+
+        axis = list(range(self.ndim))
+        axis[-1], axis[-2] = axis[-2], axis[-1]
+
+        return self.transpose(axis)
+
     def swapaxes(self, axis1, axis2):
         """Returns array that has axes axis1 and axis2 swapped.
 
@@ -1468,6 +1478,8 @@ class COO(SparseArray, NDArrayOperatorsMixin):  # lgtm [py/missing-equals]
         (array([0, 1, 2, 3, 4]), array([0, 1, 2, 3, 4]))
         """
         check_zero_fill_value(self)
+        if self.ndim == 0:
+            raise ValueError("`nonzero` is undefined for `self.ndim == 0`.")
         return tuple(self.coords)
 
     def asformat(self, format, **kwargs):
diff --git a/sparse/numba_backend/_sparse_array.py b/sparse/numba_backend/_sparse_array.py
index 0c9b2e0..763a779 100644
--- a/sparse/numba_backend/_sparse_array.py
+++ b/sparse/numba_backend/_sparse_array.py
@@ -47,6 +47,17 @@ class SparseArray:
 
     dtype = None
 
+    @property
+    def device(self):
+        data = getattr(self, "data", None)
+        return getattr(data, "device", "cpu")
+
+    def to_device(self, device, /, *, stream=None):
+        if device != "cpu":
+            raise ValueError("Only `device='cpu'` is supported.")
+
+        return self
+
     @property
     @abstractmethod
     def nnz(self):
@@ -315,11 +326,11 @@ class SparseArray:
             return self.__array_function__(ufunc, (np.ndarray, type(self)), inputs, kwargs)
 
         if out is not None:
-            test_args = [np.empty(1, dtype=a.dtype) if hasattr(a, "dtype") else [a] for a in inputs]
+            test_args = [np.empty((1,), dtype=a.dtype) if hasattr(a, "dtype") else a for a in inputs]
             test_kwargs = kwargs.copy()
             if method == "reduce":
                 test_kwargs["axis"] = None
-            test_out = tuple(np.empty(1, dtype=a.dtype) for a in out)
+            test_out = tuple(np.empty((1,), dtype=a.dtype) for a in out)
             if len(test_out) == 1:
                 test_out = test_out[0]
             getattr(ufunc, method)(*test_args, out=test_out, **test_kwargs)
@@ -379,13 +390,9 @@ class SparseArray:
         """
         axis = normalize_axis(axis, self.ndim)
         zero_reduce_result = method.reduce([self.fill_value, self.fill_value], **kwargs)
-        reduce_super_ufunc = None
-
-        if not equivalent(zero_reduce_result, self.fill_value):
-            reduce_super_ufunc = _reduce_super_ufunc.get(method)
-
-            if reduce_super_ufunc is None:
-                raise ValueError(f"Performing this reduction operation would produce a dense result: {method!s}")
+        reduce_super_ufunc = _reduce_super_ufunc.get(method)
+        if not equivalent(zero_reduce_result, self.fill_value) and reduce_super_ufunc is None:
+            raise ValueError(f"Performing this reduction operation would produce a dense result: {method!s}")
 
         if not isinstance(axis, tuple):
             axis = (axis,)
@@ -788,7 +795,7 @@ class SparseArray:
         if dtype is None and issubclass(self.dtype.type, np.integer | np.bool_):
             dtype = np.dtype("f8")
 
-        arrmean = self.sum(axis, dtype=dtype, keepdims=True)
+        arrmean = self.sum(axis, dtype=dtype, keepdims=True)[...]
         np.divide(arrmean, rcount, out=arrmean)
         x = self - arrmean
         if issubclass(self.dtype.type, np.complexfloating):
diff --git a/sparse/numba_backend/_umath.py b/sparse/numba_backend/_umath.py
index 1d77e1a..786f062 100644
--- a/sparse/numba_backend/_umath.py
+++ b/sparse/numba_backend/_umath.py
@@ -431,7 +431,7 @@ class _Elemwise:
                 processed_args.append(COO.from_scipy_sparse(arg))
             elif isscalar(arg) or isinstance(arg, np.ndarray):
                 # Faster and more reliable to pass ()-shaped ndarrays as scalars.
-                processed_args.append(np.asarray(arg))
+                processed_args.append(arg)
             elif isinstance(arg, SparseArray):
                 if not isinstance(arg, COO):
                     arg = arg.asformat(COO)
@@ -513,15 +513,22 @@ class _Elemwise:
         """
         from ._coo import COO
 
-        zero_args = tuple(
-            np.asarray(arg.fill_value, like=arg.data) if isinstance(arg, COO) else arg for arg in self.args
-        )
+        def get_zero_arg(x):
+            if isinstance(x, COO):
+                return np.atleast_1d(x.fill_value)
+
+            if isinstance(x, np.generic | np.ndarray):
+                return np.atleast_1d(x)
+
+            return x
+
+        zero_args = tuple(get_zero_arg(a) for a in self.args)
 
         # Some elemwise functions require a dtype argument, some abhorr it.
         try:
-            fill_value_array = self.func(*np.broadcast_arrays(*zero_args), dtype=self.dtype, **self.kwargs)
+            fill_value_array = self.func(*zero_args, dtype=self.dtype, **self.kwargs)
         except TypeError:
-            fill_value_array = self.func(*np.broadcast_arrays(*zero_args), **self.kwargs)
+            fill_value_array = self.func(*zero_args, **self.kwargs)
 
         try:
             fill_value = fill_value_array[(0,) * fill_value_array.ndim]
@@ -531,7 +538,7 @@ class _Elemwise:
             )
             fill_value = self.func(*zero_args, **self.kwargs)[()]
 
-        equivalent_fv = equivalent(fill_value, fill_value_array).all()
+        equivalent_fv = equivalent(fill_value, fill_value_array, loose=True).all()
         if not equivalent_fv and self.shape != self.ndarray_shape:
             raise ValueError(
                 "Performing a mixed sparse-dense operation that would result in a dense array. "
@@ -558,7 +565,7 @@ class _Elemwise:
         """
         from ._coo import COO
 
-        full_shape = _get_nary_broadcast_shape(*tuple(arg.shape for arg in self.args))
+        full_shape = _get_nary_broadcast_shape(*tuple(np.shape(arg) for arg in self.args))
         non_ndarray_shape = _get_nary_broadcast_shape(*tuple(arg.shape for arg in self.args if isinstance(arg, COO)))
         ndarray_shape = _get_nary_broadcast_shape(*tuple(arg.shape for arg in self.args if isinstance(arg, np.ndarray)))
 
@@ -587,7 +594,7 @@ class _Elemwise:
         ndarray_args = [arg for arg, m in zip(self.args, mask, strict=True) if m is None]
 
         matched_broadcast_shape = _get_nary_broadcast_shape(
-            *tuple(arg.shape for arg in itertools.chain(matched_args, ndarray_args))
+            *tuple(np.shape(arg) for arg in itertools.chain(matched_args, ndarray_args))
         )
 
         matched_arrays = self._match_coo(*matched_args, cache=self.cache, broadcast_shape=matched_broadcast_shape)
diff --git a/sparse/numba_backend/_utils.py b/sparse/numba_backend/_utils.py
index 08a39ec..8d1fb5e 100644
--- a/sparse/numba_backend/_utils.py
+++ b/sparse/numba_backend/_utils.py
@@ -406,7 +406,7 @@ def normalize_axis(axis, ndim):
     raise ValueError(f"axis {axis} not understood")
 
 
-def equivalent(x, y):
+def equivalent(x, y, /, loose=False):
     """
     Checks the equivalence of two scalars or arrays with broadcasting. Assumes
     a consistent dtype.
@@ -432,17 +432,27 @@ def equivalent(x, y):
     >>> equivalent(np.inf, np.inf)
     True
     >>> equivalent(np.PZERO, np.NZERO)
-    True
+    False
     """
     x = np.asarray(x)
     y = np.asarray(y)
     # Can't contain NaNs
-    if any(np.issubdtype(x.dtype, t) for t in [np.integer, np.bool_, np.character]):
+    dt = np.result_type(x.dtype, y.dtype)
+    if not any(np.issubdtype(dt, t) for t in [np.floating, np.complexfloating]):
         return x == y
 
-    # Can contain NaNs
-    # FIXME: Complex floats and np.void with multiple values can't be compared properly.
-    return (x == y) | ((x != x) & (y != y))  # noqa: PLR0124
+    if loose:
+        if np.issubdtype(dt, np.complexfloating):
+            return equivalent(x.real, y.real) & equivalent(x.imag, y.imag)
+
+        # TODO: Rec array handling
+        return (x == y) | ((x != x) & (y != y))
+
+    if x.size == 0 or y.size == 0:
+        shape = np.broadcast_shapes(x.shape, y.shape)
+        return np.empty(shape, dtype=np.bool_)
+    x, y = np.broadcast_arrays(x[..., None], y[..., None])
+    return (x.astype(dt).view(np.uint8) == y.astype(dt).view(np.uint8)).all(axis=-1)
 
 
 # copied from zarr
@@ -549,7 +559,7 @@ def check_fill_value(x, /, *, accept_fv=None) -> None:
     if not isinstance(accept_fv, Iterable):
         accept_fv = [accept_fv]
 
-    if not any(equivalent(fv, x.fill_value) for fv in accept_fv):
+    if not any(equivalent(fv, x.fill_value, loose=True) for fv in accept_fv):
         raise ValueError(f"{x.fill_value=} but should be in {accept_fv}.")
 
 


Test patch:
diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py
index 0c548c6..ed1ef96 100644
--- a/sparse/numba_backend/tests/test_coo.py
+++ b/sparse/numba_backend/tests/test_coo.py
@@ -1793,14 +1793,19 @@ def test_expand_dims(axis):
 @pytest.mark.parametrize("fill_value", [-1, 0, 1, 3])
 @pytest.mark.parametrize("axis", [0, 1, -1])
 @pytest.mark.parametrize("descending", [False, True])
-def test_sort(arr, fill_value, axis, descending):
+@pytest.mark.parametrize(
+    "stable", [False, pytest.param(True, marks=pytest.mark.xfail(reason="Numba doesn't support `stable=True`."))]
+)
+def test_sort(arr, fill_value, axis, descending, stable):
     if axis >= arr.ndim:
         return
 
     s_arr = sparse.COO.from_numpy(arr, fill_value)
 
-    result = sparse.sort(s_arr, axis=axis, descending=descending)
-    expected = -np.sort(-arr, axis=axis) if descending else np.sort(arr, axis=axis)
+    kind = "mergesort" if stable else "quicksort"
+
+    result = sparse.sort(s_arr, axis=axis, descending=descending, stable=stable)
+    expected = -np.sort(-arr, axis=axis, kind=kind) if descending else np.sort(arr, axis=axis, kind=kind)
 
     np.testing.assert_equal(result.todense(), expected)
     # make sure no inplace changes happened
@@ -1868,27 +1873,25 @@ def test_matrix_transpose(ndim, density):
     expected = np.transpose(xd, axes=transpose_axes)
     actual = sparse.matrix_transpose(xs)
 
-    np.testing.assert_equal(actual.todense(), expected)
+    assert_eq(actual, expected)
+    assert_eq(xs.mT, expected)
 
 
 @pytest.mark.parametrize(
-    "shape1, shape2",
+    ("shape1", "shape2", "axis"),
     [
-        ((2, 3, 4), (3, 4)),
-        ((3, 4), (2, 3, 4)),
-        ((3, 1, 4), (3, 2, 4)),
-        ((1, 3, 4), (3, 4)),
-        ((3, 4, 1), (3, 4, 2)),
-        ((1, 5), (5, 1)),
-        ((3, 1), (3, 4)),
-        ((3, 1), (1, 4)),
-        ((1, 4), (3, 4)),
-        ((2, 2, 2), (1, 1, 1)),
+        ((2, 3, 4), (3, 4), -2),
+        ((3, 4), (2, 3, 4), -1),
+        ((3, 1, 4), (3, 2, 4), 2),
+        ((1, 3, 4), (3, 4), -2),
+        ((3, 4, 1), (3, 4, 2), 0),
+        ((3, 1), (3, 4), -2),
+        ((1, 4), (3, 4), 1),
     ],
 )
 @pytest.mark.parametrize("density", [0.0, 0.1, 0.25, 1.0])
 @pytest.mark.parametrize("is_complex", [False, True])
-def test_vecdot(shape1, shape2, density, rng, is_complex):
+def test_vecdot(shape1, shape2, axis, density, rng, is_complex):
     def data_rvs(size):
         data = rng.random(size)
         if is_complex:
@@ -1898,8 +1901,6 @@ def test_vecdot(shape1, shape2, density, rng, is_complex):
     s1 = sparse.random(shape1, density=density, data_rvs=data_rvs)
     s2 = sparse.random(shape2, density=density, data_rvs=data_rvs)
 
-    axis = rng.integers(max(s1.ndim, s2.ndim))
-
     x1 = s1.todense()
     x2 = s2.todense()
 
@@ -1909,7 +1910,72 @@ def test_vecdot(shape1, shape2, density, rng, is_complex):
 
         return np.sum(x1 * x2, axis=axis)
 
-    expected = np_vecdot(x1, x2, axis=axis)
     actual = sparse.vecdot(s1, s2, axis=axis)
+    expected = np_vecdot(x1, x2, axis=axis)
 
     np.testing.assert_allclose(actual.todense(), expected)
+
+
+@pytest.mark.parametrize(
+    ("shape1", "shape2", "axis"),
+    [
+        ((2, 3, 4), (3, 4), 0),
+        ((3, 4), (2, 3, 4), 0),
+        ((3, 1, 4), (3, 2, 4), -2),
+        ((1, 3, 4), (3, 4), -3),
+        ((3, 4, 1), (3, 4, 2), -1),
+        ((3, 1), (3, 4), 1),
+        ((1, 4), (3, 4), -2),
+    ],
+)
+def test_vecdot_invalid_axis(shape1, shape2, axis):
+    s1 = sparse.random(shape1, density=0.5)
+    s2 = sparse.random(shape2, density=0.5)
+
+    with pytest.raises(ValueError, match=r"Shapes must match along"):
+        sparse.vecdot(s1, s2, axis=axis)
+
+
+@pytest.mark.parametrize(
+    ("func", "args", "kwargs"),
+    [
+        (sparse.eye, (5,), {}),
+        (sparse.zeros, ((5,)), {}),
+        (sparse.ones, ((5,)), {}),
+        (sparse.full, ((5,), 5), {}),
+        (sparse.empty, ((5,)), {}),
+        (sparse.full_like, (5,), {}),
+        (sparse.ones_like, (), {}),
+        (sparse.zeros_like, (), {}),
+        (sparse.empty_like, (), {}),
+        (sparse.asarray, (), {}),
+    ],
+)
+def test_invalid_device(func, args, kwargs):
+    if func.__name__.endswith("_like") or func is sparse.asarray:
+        like = sparse.random((5, 5), density=0.5)
+        args = (like,) + args
+
+    with pytest.raises(ValueError, match="Device must be"):
+        func(*args, device="invalid_device", **kwargs)
+
+
+def test_device():
+    s = sparse.random((5, 5), density=0.5)
+    data = getattr(s, "data", None)
+    device = getattr(data, "device", "cpu")
+
+    assert s.device == device
+
+
+def test_to_device():
+    s = sparse.random((5, 5), density=0.5)
+    s2 = s.to_device(s.device)
+
+    assert s is s2
+
+
+def test_to_invalid_device():
+    s = sparse.random((5, 5), density=0.5)
+    with pytest.raises(ValueError, match=r"Only .* is supported."):
+        s.to_device("invalid_device")
diff --git a/sparse/numba_backend/tests/test_namespace.py b/sparse/numba_backend/tests/test_namespace.py
index 902b1ff..39556f9 100644
--- a/sparse/numba_backend/tests/test_namespace.py
+++ b/sparse/numba_backend/tests/test_namespace.py
@@ -43,6 +43,7 @@ def test_namespace():
         "complex64",
         "concat",
         "concatenate",
+        "conj",
         "cos",
         "cosh",
         "diagonal",