Repository: pydata/sparse

Issue:
Dot should distinguish between vectors and matrices.
For example, if `x` is a vector and `y` is a matrix, `dot` should recognize the difference (and also return a vector). Same for the other way around, and also for two vectors.

Discussion:

Patch:
diff --git a/docs/operations.rst b/docs/operations.rst
index 0e0b5e9..7e4ef2b 100644
--- a/docs/operations.rst
+++ b/docs/operations.rst
@@ -288,6 +288,6 @@ All of the following will raise an :obj:`IndexError`, like in Numpy 1.13 and lat
 Other Operations
 ----------------
 :obj:`COO` arrays support a number of other common operations. Among them are
-:obj:`dot`, :obj:`tensordot`, :obj:`concatenate` and :obj:`stack`,
-:obj:`COO.transpose` and :obj:`COO.reshape`. You can view the full list on the
-API reference page for :obj:`sparse`
+:obj:`dot <COO.dot>`, :obj:`tensordot <COO.tensordot>`, :obj:`concatenate <COO.concatenate>`
+and :obj:`stack <COO.stack>`, :obj:`COO.transpose <COO.transpose>` and :obj:`reshape <COO.reshape>`.
+You can view the full list on the API reference page for :obj:`sparse`.
diff --git a/sparse/coo.py b/sparse/coo.py
index d85045c..71dd450 100644
--- a/sparse/coo.py
+++ b/sparse/coo.py
@@ -1738,7 +1738,17 @@ def dot(a, b):
         raise NotImplementedError(
             "Cannot perform dot product on types %s, %s" %
             (type(a), type(b)))
-    return tensordot(a, b, axes=((a.ndim - 1,), (b.ndim - 2,)))
+
+    if a.ndim == 1 and b.ndim == 1:
+        return (a * b).sum()
+
+    a_axis = -1
+    b_axis = -2
+
+    if b.ndim == 1:
+        b_axis = -1
+
+    return tensordot(a, b, axes=(a_axis, b_axis))
 
 
 def _dot(a, b):


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index ae3889c..b7fab97 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -158,10 +158,16 @@ def test_tensordot(a_shape, b_shape, axes):
     # assert isinstance(sparse.tensordot(a, sb, axes), COO)
 
 
-def test_dot():
-    import operator
-    sa = sparse.random((3, 4, 5), density=0.5)
-    sb = sparse.random((5, 6), density=0.5)
+@pytest.mark.parametrize('a_shape, b_shape', [
+    ((3, 4, 5), (5, 6)),
+    ((4, 5), (5, 6)),
+    ((5,), (5, 6)),
+    ((4, 5), (5,)),
+    ((5,), (5,)),
+])
+def test_dot(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()