BUG: Fix CSR/CSC matmul (#660)

Author: mtsokol

sparse/pydata_backend/_compressed/compressed.py

@@ -2,6 +2,7 @@
 import operator
 from collections.abc import Iterable
 from functools import reduce
+from typing import Union
 
 import numpy as np
 from numpy.lib.mixins import NDArrayOperatorsMixin
@@ -876,11 +877,14 @@ def from_scipy_sparse(cls, x):
         x = x.asformat("csr", copy=False)
         return cls((x.data, x.indices, x.indptr), shape=x.shape)
 
-    def transpose(self, axes: None = None, copy: bool = False) -> "CSC":
-        if axes is not None:
-            raise ValueError
+    def transpose(self, axes: None = None, copy: bool = False) -> Union["CSC", "CSR"]:
+        axes = normalize_axis(axes, self.ndim)
+        if axes not in [(0, 1), (1, 0), None]:
+            raise ValueError(f"Invalid transpose axes: {axes}")
         if copy:
             self = self.copy()
+        if axes == (0, 1):
+            return self
         return CSC((self.data, self.indices, self.indptr), self.shape[::-1])
 
 
@@ -905,9 +909,12 @@ def from_scipy_sparse(cls, x):
         x = x.asformat("csc", copy=False)
         return cls((x.data, x.indices, x.indptr), shape=x.shape)
 
-    def transpose(self, axes: None = None, copy: bool = False) -> CSR:
-        if axes is not None:
-            raise ValueError
+    def transpose(self, axes: None = None, copy: bool = False) -> Union["CSC", "CSR"]:
+        axes = normalize_axis(axes, self.ndim)
+        if axes not in [(0, 1), (1, 0), None]:
+            raise ValueError(f"Invalid transpose axes: {axes}")
         if copy:
             self = self.copy()
+        if axes == (0, 1):
+            return self
         return CSR((self.data, self.indices, self.indptr), self.shape[::-1])

sparse/pydata_backend/tests/test_compressed_2d.py

@@ -41,7 +41,7 @@ def data_rvs(n):
 
     else:
         data_rvs = None
-    return cls(sparse.random((20, 30, 40), density=0.25, data_rvs=data_rvs).astype(dtype))
+    return cls(sparse.random((20, 20), density=0.25, data_rvs=data_rvs).astype(dtype))
 
 
 def test_repr(random_sparse):
@@ -111,7 +111,21 @@ def test_transpose(random_sparse, copy):
     assert_eq(random_sparse, tt)
     assert type(random_sparse) == type(tt)
 
+    assert_eq(random_sparse.transpose(axes=(0, 1)), random_sparse)
+    assert_eq(random_sparse.transpose(axes=(1, 0)), t)
+    with pytest.raises(ValueError, match="Invalid transpose axes"):
+        random_sparse.transpose(axes=0)
+
 
 def test_transpose_error(random_sparse):
     with pytest.raises(ValueError):
         random_sparse.transpose(axes=1)
+
+
+def test_matmul(random_sparse_small):
+    arr = random_sparse_small.todense()
+
+    actual = random_sparse_small @ random_sparse_small
+    expected = arr @ arr
+
+    assert_eq(actual, expected)