BUG: fix up matmult/dot

Author: ev-br

torch_np/_detail/_ufunc_impl.py

@@ -1,6 +1,7 @@
 import torch
 
 from . import _util
+from . import _dtypes_impl
 
 
 def deco_ufunc(torch_func):
@@ -70,7 +71,6 @@ def wrapped(
 logical_and = deco_ufunc(torch.logical_and)
 logical_or = deco_ufunc(torch.logical_or)
 logical_xor = deco_ufunc(torch.logical_xor)
-matmul = deco_ufunc(torch.matmul)
 maximum = deco_ufunc(torch.maximum)
 minimum = deco_ufunc(torch.minimum)
 remainder = deco_ufunc(torch.remainder)
@@ -143,7 +143,15 @@ def _absolute(x):
         return x
     return torch.absolute(x)
 
+def _matmul(x, y):
+    # work around RuntimeError: expected scalar type Int but found Double
+    dtype = _dtypes_impl.result_type_impl((x.dtype, y.dtype))
+    x = x.to(dtype)
+    y = y.to(dtype)
+    result = torch.matmul(x, y)
+    return result
 
 cbrt = deco_ufunc(_cbrt)
 positive = deco_ufunc(_positive)
 absolute = deco_ufunc(_absolute)
+matmul = deco_ufunc(_matmul)

torch_np/_detail/implementations.py

@@ -178,6 +178,8 @@ def trace(tensor, offset=0, axis1=0, axis2=1, dtype=None, out=None):
 def diagonal(tensor, offset=0, axis1=0, axis2=1):
     axis1 = _util.normalize_axis_index(axis1, tensor.ndim)
     axis2 = _util.normalize_axis_index(axis2, tensor.ndim)
+    if axis1 == axis2:
+        raise ValueError("axis1 and axis2 cannot be the same")
     result = torch.diagonal(tensor, offset, axis1, axis2)
     return result
 
@@ -492,17 +494,25 @@ def arange(start=None, stop=None, step=1, dtype=None):
     if start is None:
         start = 0
 
-    if dtype is None:
-        dt_list = [_util._coerce_to_tensor(x).dtype for x in (start, stop, step)]
-        dtype = _dtypes_impl.default_int_dtype
-        dt_list.append(dtype)
-        dtype = _dtypes_impl.result_type_impl(dt_list)
+#    if dtype is None:
+    dt_list = [_util._coerce_to_tensor(x).dtype for x in (start, stop, step)]
+    dtype = _dtypes_impl.default_int_dtype
+    dt_list.append(dtype)
+    dtype = _dtypes_impl.result_type_impl(dt_list)
 
+    # work around RuntimeError: "arange_cpu" not implemented for 'ComplexFloat'        
+    orig_dtype = dtype
+    is_complex = dtype is not None and dtype.is_complex
     try:
-        return torch.arange(start, stop, step, dtype=dtype)
+        if is_complex:
+            dtype = torch.float64
+        result = torch.arange(start, stop, step, dtype=orig_dtype)
+        if is_complex:
+            result  = result.to(dttype)
     except RuntimeError:
         raise ValueError("Maximum allowed size exceeded")
 
+    return result
 
 # ### empty/full et al ###
 

torch_np/_funcs.py

@@ -1,7 +1,7 @@
 import torch
 
 from . import _decorators, _helpers
-from ._detail import _flips, _util
+from ._detail import _flips, _util, _dtypes_impl
 from ._detail import implementations as _impl
 
 
@@ -101,6 +101,9 @@ def vdot(a, b, /):
 
 def dot(a, b, out=None):
     t_a, t_b = _helpers.to_tensors(a, b)
+    dtype = _dtypes_impl.result_type_impl((t_a.dtype, t_b.dtype))
+    t_a = t_a.to(dtype)
+    t_b = t_b.to(dtype)
     result = _impl.dot(t_a, t_b)
     return _helpers.result_or_out(result, out)
 

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -2468,7 +2468,6 @@ def test_arr_mult_2(self, func):
                 func(edf, edf[:, ::-1].T.copy())
             )
 
-    @pytest.mark.xfail(reason="TODO np.dot")
     @pytest.mark.parametrize('func', (np.dot, np.matmul))
     @pytest.mark.parametrize('dtype', 'ifdFD')
     def test_no_dgemv(self, func, dtype):