Repository: pydata/sparse

Issue:
COO.__matmul__ gives the wrong shape of resultant array
Please see the following, COO.__matmul__ is not consistent with np.ndarray.__matmul__
```python
In [1]: import numpy as np

In [2]: import sparse

In [3]: npx = np.random.random((1,2,3,4))

In [4]: npy = np.random.random((1,2,4,3))

In [5]: npz = npx @ npy

In [6]: npz.shape
Out[6]: (1, 2, 3, 3)

In [7]: spx = sparse.random((1,2,3,4), 0.8)

In [8]: spy = sparse.random((1,2,4,3), 0.8)

In [9]: spz = spx @ spy

In [10]: type(spz)
Out[10]: numpy.ndarray

In [11]: spz.shape
Out[11]: (1, 2, 3, 1, 2, 3)
```

Discussion:
gongliyu: Sure, I will add some test cases. Do I need to close this PR and open another one after I add the tests in my forked repo? Does pydata/sparse contains a script to run the tests before push to remote repo? I'm sorry I'm very new to the github PR feature. I will really appreciate it if you can help me on this.

Thanks!
hameerabbasi: Hi! Like you, I was also very new a year or so ago! I'd like to personally thank (and congratulate) you for doing this.

Usually it's recommended to push your changes to a branch and keep your `master` branch sync'd to the main branch, but for this PR it's okay.

If you push changes to your branch, this PR will be automatically updated.
hameerabbasi: To run the tests, you can run `pip install -e .[all]` in the root of the repository and then run `py.test` there. It will run all tests for you, including for code style.
hameerabbasi: This pull request **introduces 1 alert** when merging 99c15a904a743501baf5b8ce51a122855d6df8c8 into 87f414c0e63c3c3213225d2db0d354ea51b4f3f2 - [view on LGTM.com](https://lgtm.com/projects/g/pydata/sparse/rev/pr-e46ffd4c0a5c90ec43636e43a8d714a5e6e3cbbb)

**new alerts:**

* 1 for Syntax error

---

*Comment posted by [LGTM.com](https://lgtm.com)*
gongliyu: I realized that np.dot and np.matmul are different for some cases, 
https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.matmul.html
So I add a function: common.matmul with some tests, and redirect COO.__matmul__ to common.matmul. 
hameerabbasi: This pull request **introduces 1 alert** when merging 6fe10d51e7020fdfffe7bfde8d8fc27f5d05de14 into 87f414c0e63c3c3213225d2db0d354ea51b4f3f2 - [view on LGTM.com](https://lgtm.com/projects/g/pydata/sparse/rev/pr-893a662ddb8f544d20edc38c907cd7f91ce336af)

**new alerts:**

* 1 for Syntax error

---

*Comment posted by [LGTM.com](https://lgtm.com)*
gongliyu: I tried to read the report and improve it. If it still can not pass the coverage test, please help me add some test cases.
hameerabbasi: Sure thing! You will need to do one more thing, add docs for `matmul`. You need to edit the file `docs/generated/sparse.rst` and add `matmul` (everything is in alphabetical order). Then run `sphinx-build -W -b html . _build/html` while in the `docs/` directory.

Patch:
diff --git a/docs/generated/sparse.matmul.rst b/docs/generated/sparse.matmul.rst
new file mode 100644
index 0000000..4db8012
--- /dev/null
+++ b/docs/generated/sparse.matmul.rst
@@ -0,0 +1,6 @@
+matmul
+======
+
+.. currentmodule:: sparse
+
+.. autofunction:: matmul
\ No newline at end of file
diff --git a/docs/generated/sparse.rst b/docs/generated/sparse.rst
index 9551e5a..1d78385 100644
--- a/docs/generated/sparse.rst
+++ b/docs/generated/sparse.rst
@@ -41,6 +41,8 @@ API
 
     load_npz
 
+    matmul
+
     nanmax
 
     nanmean
diff --git a/sparse/coo/__init__.py b/sparse/coo/__init__.py
index 0d4ef10..b4ea805 100644
--- a/sparse/coo/__init__.py
+++ b/sparse/coo/__init__.py
@@ -1,6 +1,6 @@
 from .core import COO, as_coo
 from .umath import elemwise
-from .common import (tensordot, dot, concatenate, stack, triu, tril, where,
+from .common import (tensordot, dot, matmul, concatenate, stack, triu, tril, where,
                      nansum, nanmean, nanprod, nanmin, nanmax, nanreduce, roll,
                      eye, full, full_like, zeros, zeros_like, ones, ones_like,
                      kron)
diff --git a/sparse/coo/common.py b/sparse/coo/common.py
index c6d022d..0187c9a 100644
--- a/sparse/coo/common.py
+++ b/sparse/coo/common.py
@@ -86,6 +86,11 @@ def tensordot(a, b, axes=2):
     from .core import COO
     check_zero_fill_value(a, b)
 
+    if scipy.sparse.issparse(a):
+        a = asCOO(a)
+    if scipy.sparse.issparse(b):
+        b = asCOO(b)
+
     try:
         iter(axes)
     except TypeError:
@@ -154,6 +159,76 @@ def tensordot(a, b, axes=2):
     return res.reshape(olda + oldb)
 
 
+def matmul(a, b):
+    """Perform the equivalent of :obj:`numpy.matmul` on two arrays.
+
+    Parameters
+    ----------
+    a, b : Union[COO, np.ndarray, scipy.sparse.spmatrix]
+        The arrays to perform the :code:`matmul` operation on.
+
+    Returns
+    -------
+    Union[COO, numpy.ndarray]
+        The result of the operation.
+
+    Raises
+    ------
+    ValueError
+        If all arguments don't have zero fill-values, or the shape of the two arrays is not broadcastable.
+
+    See Also
+    --------
+    numpy.matmul : NumPy equivalent function.
+    COO.__matmul__ : Equivalent function for COO objects.
+    """
+    check_zero_fill_value(a, b)
+    if not hasattr(a, 'ndim') or not hasattr(b, 'ndim'):
+        raise TypeError(
+            "Cannot perform dot product on types %s, %s" %
+            (type(a), type(b)))
+
+    # When one of the input is less than 2-d, it is equivalent to dot
+    if a.ndim <= 2 or b.ndim <= 2:
+        return dot(a, b)
+
+    # If a can be squeeze to a vector, use dot will be faster
+    if a.ndim <= b.ndim and np.prod(a.shape[:-1]) == 1:
+        res = dot(a.reshape(-1), b)
+        shape = list(res.shape)
+        shape.insert(-1, 1)
+        return res.reshape(shape)
+
+    # If b can be squeeze to a matrix, use dot will be faster
+    if b.ndim <= a.ndim and np.prod(b.shape[:-2]) == 1:
+        return dot(a, b.reshape(b.shape[-2:]))
+
+    if a.ndim < b.ndim:
+        a = a[(None,) * (b.ndim - a.ndim)]
+    if a.ndim > b.ndim:
+        b = b[(None,) * (a.ndim - b.ndim)]
+    for i, j in zip(a.shape[:-2], b.shape[:-2]):
+        if i != 1 and j != 1 and i != j:
+            raise ValueError('shapes of a and b are not broadcastable')
+
+    def _matmul_recurser(a, b):
+        if a.ndim == 2:
+            return dot(a, b)
+        res = []
+        for i in range(max(a.shape[0], b.shape[0])):
+            a_i = a[0] if a.shape[0] == 1 else a[i]
+            b_i = b[0] if b.shape[0] == 1 else b[i]
+            res.append(_matmul_recurser(a_i, b_i))
+        mask = [isinstance(x, SparseArray) for x in res]
+        if all(mask):
+            return stack(res)
+        else:
+            res = [x.todense() if isinstance(x, SparseArray) else x
+                   for x in res]
+            return np.stack(res)
+    return _matmul_recurser(a, b)
+
+
 def dot(a, b):
     """
     Perform the equivalent of :obj:`numpy.dot` on two arrays.
@@ -192,7 +267,6 @@ def dot(a, b):
 
     if b.ndim == 1:
         b_axis = -1
-
     return tensordot(a, b, axes=(a_axis, b_axis))
 
 
@@ -201,11 +275,13 @@ def _dot(a, b):
 
     if isinstance(b, COO) and not isinstance(a, COO):
         return _dot(b.T, a.T).T
-    aa = a.tocsr()
+
+    if isinstance(a, (COO, scipy.sparse.spmatrix)):
+        a = a.tocsr()
 
     if isinstance(b, (COO, scipy.sparse.spmatrix)):
         b = b.tocsc()
-    return aa.dot(b)
+    return a.dot(b)
 
 
 def kron(a, b):
diff --git a/sparse/coo/core.py b/sparse/coo/core.py
index de0524a..9b97029 100644
--- a/sparse/coo/core.py
+++ b/sparse/coo/core.py
@@ -7,7 +7,7 @@ import numpy as np
 import scipy.sparse
 from numpy.lib.mixins import NDArrayOperatorsMixin
 
-from .common import dot
+from .common import dot, matmul
 from .indexing import getitem
 from .umath import elemwise, broadcast_to
 from ..compatibility import int, range
@@ -1405,13 +1405,13 @@ class COO(SparseArray, NDArrayOperatorsMixin):
 
     def __matmul__(self, other):
         try:
-            return dot(self, other)
+            return matmul(self, other)
         except NotImplementedError:
             return NotImplemented
 
     def __rmatmul__(self, other):
         try:
-            return dot(other, self)
+            return matmul(other, self)
         except NotImplementedError:
             return NotImplemented
 


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index c403b3d..c56cea5 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -272,6 +272,49 @@ def test_tensordot(a_shape, b_shape, axes):
 
 
 @pytest.mark.parametrize('a_shape, b_shape', [
+    ((3, 1, 6, 5), (2, 1, 4, 5, 6)),
+    ((2, 1, 4, 5, 6), (3, 1, 6, 5)),
+    ((1, 1, 5), (3, 5, 6)),
+    ((3, 4, 5), (1, 5, 6)),
+    ((3, 4, 5), (3, 5, 6)),
+    ((3, 4, 5), (5, 6)),
+    ((4, 5), (5, 6)),
+    ((5,), (5, 6)),
+    ((4, 5), (5,)),
+    ((5,), (5,)),
+])
+def test_matmul(a_shape, b_shape):
+    sa = sparse.random(a_shape, density=0.5)
+    sb = sparse.random(b_shape, density=0.5)
+
+    a = sa.todense()
+    b = sb.todense()
+
+    assert_eq(np.matmul(a, b), sparse.matmul(sa, sb))
+    assert_eq(sparse.matmul(sa, b), sparse.matmul(a, sb))
+    assert_eq(sparse.matmul(a, b), sparse.matmul(sa, sb))
+
+    if a.ndim == 2 or b.ndim == 2:
+        assert_eq(
+            np.matmul(a, b),
+            sparse.matmul(scipy.sparse.coo_matrix(a) if a.ndim == 2 else sa,
+                          scipy.sparse.coo_matrix(b) if b.ndim == 2 else sb))
+
+    if hasattr(operator, 'matmul'):
+        assert_eq(operator.matmul(a, b), operator.matmul(sa, sb))
+
+
+def test_matmul_errors():
+    with pytest.raises(ValueError):
+        sa = sparse.random((3, 4, 5, 6), 0.5)
+        sb = sparse.random((3, 6, 5, 6), 0.5)
+        sparse.matmul(sa, sb)
+
+
+@pytest.mark.parametrize('a_shape, b_shape', [
+    ((1, 4, 5), (3, 5, 6)),
+    ((3, 4, 5), (1, 5, 6)),
+    ((3, 4, 5), (3, 5, 6)),
     ((3, 4, 5), (5, 6)),
     ((4, 5), (5, 6)),
     ((5,), (5, 6)),
@@ -287,11 +330,12 @@ def test_dot(a_shape, b_shape):
 
     assert_eq(a.dot(b), sa.dot(sb))
     assert_eq(np.dot(a, b), sparse.dot(sa, sb))
+    assert_eq(sparse.dot(sa, b), sparse.dot(a, sb))
+    assert_eq(sparse.dot(a, b), sparse.dot(sa, sb))
 
     if hasattr(operator, 'matmul'):
         # Basic equivalences
         assert_eq(operator.matmul(a, b), operator.matmul(sa, sb))
-
         # Test that SOO's and np.array's combine correctly
         # Not possible due to https://github.com/numpy/numpy/issues/9028
         # assert_eq(eval("a @ sb"), eval("sa @ b"))