TST: inner, vdot

Author: ev-br

torch_np/_detail/implementations.py

@@ -692,11 +692,12 @@ def where(condition, x, y):
 
 # ### dot and other linalg ###
 
+
 def inner(t_a, t_b):
-    is_half = t_a.dtype == torch.float16 or t_b.dtype == torch.float16
-    is_bool = t_a.dtype == torch.bool or t_b.dtype == torch.bool
+    dtype = _dtypes_impl.result_type_impl((t_a.dtype, t_b.dtype))
+    is_half = dtype == torch.float16
+    is_bool = dtype == torch.bool
 
-    dtype = None
     if is_half:
         # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
         dtype = torch.float32
@@ -718,11 +719,32 @@ def inner(t_a, t_b):
 
 
 def vdot(t_a, t_b, /):
-    # torch only accepts 1D arrays, numpy ravels
+    # 1. torch only accepts 1D arrays, numpy ravels
+    # 2. torch requires matching dtype, while numpy casts (?)
     t_a, t_b = torch.atleast_1d(t_a, t_b)
     if t_a.ndim > 1:
         t_a = t_a.ravel()
     if t_b.ndim > 1:
         t_b = t_b.ravel()
+
+    dtype = _dtypes_impl.result_type_impl((t_a.dtype, t_b.dtype))
+    is_half = dtype == torch.float16
+    is_bool = dtype == torch.bool
+
+    # work around torch's "dot" not implemented for 'Half', 'Bool'
+    if is_half:
+        dtype = torch.float32
+    if is_bool:
+        dtype = torch.uint8
+
+    t_a = _util.cast_if_needed(t_a, dtype)
+    t_b = _util.cast_if_needed(t_b, dtype)
+
     result = torch.vdot(t_a, t_b)
+
+    if is_half:
+        result = result.to(torch.float16)
+    if is_bool:
+        result = result.to(torch.bool)
+
     return result

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -5421,23 +5421,22 @@ def __array_finalize__(self, obj):
         assert_(res.info == dat.info)
 
 
-@pytest.mark.xfail(reason='TODO')
 class TestVdot:
     def test_basic(self):
         dt_numeric = np.typecodes['AllFloat'] + np.typecodes['AllInteger']
         dt_complex = np.typecodes['Complex']
 
         # test real
         a = np.eye(3)
-        for dt in dt_numeric + 'O':
+        for dt in dt_numeric:
             b = a.astype(dt)
             res = np.vdot(b, b)
             assert_(np.isscalar(res))
             assert_equal(np.vdot(b, b), 3)
 
         # test complex
         a = np.eye(3) * 1j
-        for dt in dt_complex + 'O':
+        for dt in dt_complex:
             b = a.astype(dt)
             res = np.vdot(b, b)
             assert_(np.isscalar(res))
@@ -5449,6 +5448,7 @@ def test_basic(self):
         assert_(np.isscalar(res))
         assert_equal(np.vdot(b, b), True)
 
+    @pytest.mark.xfail(reason="implement order='F'")
     def test_vdot_array_order(self):
         a = np.array([[1, 2], [3, 4]], order='C')
         b = np.array([[1, 2], [3, 4]], order='F')
@@ -5476,6 +5476,20 @@ def test_vdot_uncontiguous(self):
                          np.vdot(a.flatten(), b.flatten()))
             assert_equal(np.vdot(a.copy(), b),
                          np.vdot(a.flatten(), b.flatten()))
+
+    @pytest.mark.xfail(reason="implement order='F'")
+    def test_vdot_uncontiguous_2(self):
+        # test order='F' separately
+        for size in [2, 1000]:
+            # Different sizes match different branches in vdot.
+            a = np.zeros((size, 2, 2))
+            b = np.zeros((size, 2, 2))
+            a[:, 0, 0] = np.arange(size)
+            b[:, 0, 0] = np.arange(size) + 1
+            # Make a and b uncontiguous:
+            a = a[..., 0]
+            b = b[..., 0]
+
             assert_equal(np.vdot(a.copy('F'), b),
                          np.vdot(a.flatten(), b.flatten()))
             assert_equal(np.vdot(a, b.copy('F')),