Add missing linear algebra functions from array API (#633)

* Add matrix_transpose.

* Add vecdot.

Author: hameerabbasi

docs/generated/sparse.rst

@@ -96,6 +96,8 @@ API
 
     matmul
 
+    matrix_transpose
+
     max
 
     mean
@@ -168,6 +170,8 @@ API
 
     var
 
+    vecdot
+
     where
 
     zeros

sparse/__init__.py

@@ -117,6 +117,7 @@
     sum,
     tensordot,
     var,
+    vecdot,
     zeros,
     zeros_like,
 )
@@ -135,6 +136,7 @@
     isneginf,
     isposinf,
     kron,
+    matrix_transpose,
     nanmax,
     nanmean,
     nanmin,
@@ -250,6 +252,7 @@
     "logical_or",
     "logical_xor",
     "matmul",
+    "matrix_transpose",
     "max",
     "mean",
     "min",
@@ -307,6 +310,7 @@
     "unique_counts",
     "unique_values",
     "var",
+    "vecdot",
     "where",
     "zeros",
     "zeros_like",

sparse/_common.py

@@ -2037,6 +2037,7 @@ def format_to_string(format):
 def asarray(obj, /, *, dtype=None, format="coo", backend="pydata", device=None, copy=False):
     """
     Convert the input to a sparse array.
+
     Parameters
     ----------
     obj : array_like
@@ -2051,10 +2052,12 @@ def asarray(obj, /, *, dtype=None, format="coo", backend="pydata", device=None,
         Device on which to place the created array.
     copy : bool, optional
         Boolean indicating whether or not to copy the input.
+
     Returns
     -------
     out : Union[SparseArray, numpy.ndarray]
         Sparse or 0-D array containing the data from `obj`.
+
     Examples
     --------
     >>> x = np.eye(8, dtype="i8")
@@ -2209,3 +2212,25 @@ def isfinite(x, /):
 
 def nonzero(x, /):
     return x.nonzero()
+
+
+def vecdot(x1, x2, /, *, axis=-1):
+    """
+    Computes the (vector) dot product of two arrays.
+
+    Parameters
+    ----------
+    x1, x2 : array_like
+        Input sparse arrays
+    axis : int
+        The axis to reduce over.
+
+    Returns
+    -------
+    out : Union[SparseArray, numpy.ndarray]
+        Sparse or 0-D array containing dot product.
+    """
+    if np.issubdtype(x1.dtype, np.complexfloating):
+        x1 = np.conjugate(x1)
+
+    return np.sum(x1 * x2, axis=axis)

sparse/_coo/common.py

@@ -1332,7 +1332,6 @@ def take(x, indices, /, *, axis=None):
     ------
     ValueError
         If the input array isn't and can't be converted to COO format.
-
     """
 
     x = _validate_coo_input(x)
@@ -1533,3 +1532,31 @@ def _arg_minmax_common(
         result = result.reshape([1 for _ in range(axis_none_original_ndim)])
 
     return result if keepdims else result.squeeze()
+
+
+def matrix_transpose(x, /):
+    """
+    Transposes a matrix or a stack of matrices.
+
+    Parameters
+    ----------
+    x : SparseArray
+        Input array.
+
+    Returns
+    -------
+    out : COO
+        Transposed COO array.
+
+    Raises
+    ------
+    ValueError
+        If the input array isn't and can't be converted to COO format, or if ``x.ndim < 2``.
+    """
+    if hasattr(x, "ndim") and x.ndim < 2:
+        raise ValueError("`x.ndim >= 2` must hold.")
+    x = _validate_coo_input(x)
+    transpose_axes = list(range(x.ndim))
+    transpose_axes[-2:] = transpose_axes[-2:][::-1]
+
+    return x.transpose(transpose_axes)

sparse/tests/test_coo.py

@@ -1775,7 +1775,7 @@ def test_unique_values(self, arr, fill_value):
 
     @pytest.mark.parametrize("func", [sparse.unique_counts, sparse.unique_values])
     def test_input_validation(self, func):
-        with pytest.raises(ValueError, match=r"Input must be an instance of SparseArray"):
+        with pytest.raises(ValueError, match="Input must be an instance of SparseArray"):
             func(self.arr)
 
 
@@ -1861,3 +1861,63 @@ def test_take(fill_value, indices, axis):
     expected = np.take(arr, indices, axis)
 
     np.testing.assert_equal(result.todense(), expected)
+
+
+@pytest.mark.parametrize("ndim", [2, 3, 4, 5])
+@pytest.mark.parametrize("density", [0.0, 0.1, 0.25, 1.0])
+def test_matrix_transpose(ndim, density):
+    shape = tuple(range(2, 34)[:ndim])
+    xs = sparse.random(shape, density=density)
+    xd = xs.todense()
+
+    transpose_axes = list(range(ndim))
+    transpose_axes[-2:] = transpose_axes[-2:][::-1]
+
+    expected = np.transpose(xd, axes=transpose_axes)
+    actual = sparse.matrix_transpose(xs)
+
+    np.testing.assert_equal(actual.todense(), expected)
+
+
+@pytest.mark.parametrize(
+    "shape1, shape2",
+    [
+        ((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)),
+    ],
+)
+@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 data_rvs(size):
+        data = rng.random(size)
+        if is_complex:
+            data = data + rng.random(size) * 1j
+        return data
+
+    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()
+
+    def np_vecdot(x1, x2, /, *, axis=-1):
+        if np.issubdtype(x1.dtype, np.complexfloating):
+            x1 = np.conjugate(x1)
+
+        return np.sum(x1 * x2, axis=axis)
+
+    expected = np_vecdot(x1, x2, axis=axis)
+    actual = sparse.vecdot(s1, s2, axis=axis)
+
+    np.testing.assert_allclose(actual.todense(), expected)

sparse/tests/test_namespace.py

@@ -94,6 +94,7 @@ def test_namespace():
         "logical_not",
         "logical_or",
         "logical_xor",
+        "matrix_transpose",
         "matmul",
         "max",
         "mean",
@@ -152,6 +153,7 @@ def test_namespace():
         "unique_counts",
         "unique_values",
         "var",
+        "vecdot",
         "where",
         "zeros",
         "zeros_like",