ENH: add tensordot

Author: ev-br

autogen/numpy_api_dump.py

@@ -134,10 +134,6 @@ def binary_repr(num, width=None):
     raise NotImplementedError
 
 
-def blackman(M):
-    raise NotImplementedError
-
-
 def block(arrays):
     raise NotImplementedError
 
@@ -337,14 +333,6 @@ def gradient(f, *varargs, axis=None, edge_order=1):
     raise NotImplementedError
 
 
-def hamming(M):
-    raise NotImplementedError
-
-
-def hanning(M):
-    raise NotImplementedError
-
-
 def histogram(a, bins=10, range=None, normed=None, weights=None, density=None):
     raise NotImplementedError
 
@@ -409,10 +397,6 @@ def ix_(*args):
     raise NotImplementedError
 
 
-def kaiser(M, beta):
-    raise NotImplementedError
-
-
 def lexsort(keys, axis=-1):
     raise NotImplementedError
 
@@ -759,10 +743,6 @@ def take(a, indices, axis=None, out=None, mode="raise"):
     raise NotImplementedError
 
 
-def tensordot(a, b, axes=2):
-    raise NotImplementedError
-
-
 def trapz(y, x=None, dx=1.0, axis=-1):
     raise NotImplementedError
 

torch_np/_funcs.py

@@ -1223,6 +1223,15 @@ def vdot(a: ArrayLike, b: ArrayLike, /):
     return result.item()
 
 
+@normalizer
+def tensordot(a: ArrayLike, b: ArrayLike, axes=2):
+    if isinstance(axes, (list, tuple)):
+        axes = [[ax] if isinstance(ax, int) else ax for ax in axes]
+    result = torch.tensordot(a, b, dims=axes)
+
+    return result
+
+
 @normalizer
 def dot(a: ArrayLike, b: ArrayLike, out: Optional[NDArray] = None) -> OutArray:
     dtype = _dtypes_impl.result_type_impl((a.dtype, b.dtype))

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -2988,15 +2988,21 @@ def test_shape_mismatch_error_message(self):
             np.broadcast([[1, 2, 3]], [[4], [5]], [6, 7])
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestTensordot:
 
     def test_zero_dimension(self):
         # Test resolution to issue #5663
-        a = np.ndarray((3,0))
-        b = np.ndarray((0,4))
+        a = np.zeros((3,0))
+        b = np.zeros((0,4))
         td = np.tensordot(a, b, (1, 0))
         assert_array_equal(td, np.dot(a, b))
+
+    @pytest.mark.xfail(reason="no einsum")
+    def test_zero_dimension_einsum(self):
+        # Test resolution to issue #5663
+        a = np.zeros((3,0))
+        b = np.zeros((0,4))
+        td = np.tensordot(a, b, (1, 0))
         assert_array_equal(td, np.einsum('ij,jk', a, b))
 
     def test_zero_dimensional(self):