ENH: add ndarray.__matmul__ (also __imatmul__)

Note that NumPy does not support in-place __imatmul__, but PyTorch does. So do we then.

Author: ev-br

torch_np/_ndarray.py

@@ -313,6 +313,14 @@ def __ilshift__(self, other):
     def __irshift__(self, other):
         return _binary_ufuncs.right_shift(self, other, out=self)
 
+    __matmul__ = _binary_ufuncs.matmul
+
+    def __rmatmul__(self, other):
+        return _binary_ufuncs.matmul(other, self)
+
+    def __imatmul__(self, other):
+        return _binary_ufuncs.matmul(self, other, out=self)
+
     # unary ops
     __invert__ = _unary_ufuncs.invert
     __abs__ = _unary_ufuncs.absolute

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -2605,7 +2605,6 @@ def test_diagonal_memleak(self):
         if HAS_REFCOUNT:
             assert_(sys.getrefcount(a) < 50)
 
-    @pytest.mark.xfail(reason="TODO: implement np.dot")
     def test_size_zero_memleak(self):
         # Regression test for issue 9615
         # Exercises a special-case code path for dot products of length
@@ -6071,11 +6070,11 @@ def test_matmul_bool(self):
         assert not np.any(c)
 
 
-@pytest.mark.xfail(reason='TODO: @')
 class TestMatmulOperator(MatmulCommon):
     import operator
     matmul = operator.matmul
 
+    @pytest.mark.skip(reason="no __array_priority__")
     def test_array_priority_override(self):
 
         class A:
@@ -6093,11 +6092,10 @@ def __rmatmul__(self, other):
         assert_equal(self.matmul(b, a), "A")
 
     def test_matmul_raises(self):
-        assert_raises(TypeError, self.matmul, np.int8(5), np.int8(5))
-        assert_raises(TypeError, self.matmul, np.void(b'abc'), np.void(b'abc'))
-        assert_raises(TypeError, self.matmul, np.arange(10), np.void(b'abc'))
+        assert_raises((RuntimeError, TypeError), self.matmul, np.int8(5), np.int8(5))
 
-@pytest.mark.xfail(reason='TODO @')
+
+@pytest.mark.xfail(reason="torch supports inplace matmul, and so do we")
 def test_matmul_inplace():
     # It would be nice to support in-place matmul eventually, but for now
     # we don't have a working implementation, so better just to error out