MAINT: address review comments

Author: ev-br

torch_np/_detail/_reductions.py

@@ -145,6 +145,7 @@ def mean(tensor, axis=None, dtype=None, *, where=NoValue):
 
     is_half = dtype == torch.float16
     if is_half:
+        # XXX revisit when the pytorch version has pytorch/pytorch#95166
         dtype = torch.float32
 
     if axis is None:

torch_np/_detail/implementations.py

@@ -179,8 +179,6 @@ 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
 

torch_np/_ndarray.py

@@ -371,8 +371,7 @@ def ravel(self, order="C"):
     def flatten(self, order="C"):
         if order != "C":
             raise NotImplementedError
-        # return a copy
-        result = self._tensor.ravel().clone()
+        result = self._tensor.flatten()
         return asarray(result)
 
     def nonzero(self):

torch_np/random.py

@@ -29,8 +29,8 @@
 ]
 
 
-def array_or_scalar(values, py_type=float, size=None):
-    if size is None:
+def array_or_scalar(values, py_type=float, return_scalar=False):
+    if return_scalar:
         return py_type(values.item())
     else:
         return asarray(values)
@@ -45,7 +45,7 @@ def random_sample(size=None):
     if size is None:
         size = ()
     values = torch.empty(size, dtype=_default_dtype).uniform_()
-    return array_or_scalar(values, size=size)
+    return array_or_scalar(values, return_scalar=size is None)
 
 
 def rand(*size):
@@ -60,19 +60,19 @@ def uniform(low=0.0, high=1.0, size=None):
     if size is None:
         size = ()
     values = torch.empty(size, dtype=_default_dtype).uniform_(low, high)
-    return array_or_scalar(values, size=size)
+    return array_or_scalar(values, return_scalar=size is None)
 
 
 def randn(*size):
     values = torch.randn(size, dtype=_default_dtype)
-    return array_or_scalar(values, size=size)
+    return array_or_scalar(values, return_scalar=size is None)
 
 
 def normal(loc=0.0, scale=1.0, size=None):
     if size is None:
         size = ()
     values = torch.empty(size, dtype=_default_dtype).normal_(loc, scale)
-    return array_or_scalar(values, size=size)
+    return array_or_scalar(values, return_scalar=size is None)
 
 
 def shuffle(x):
@@ -90,7 +90,7 @@ def randint(low, high=None, size=None):
     if high is None:
         low, high = 0, low
     values = torch.randint(low, high, size=size)
-    return array_or_scalar(values, int, size=size)
+    return array_or_scalar(values, int, return_scalar=size is None)
 
 
 def choice(a, size=None, replace=True, p=None):

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -1686,7 +1686,7 @@ def assert_c(arr):
         assert_fortran(a.copy('F'))
         assert_c(a.copy('A'))
 
-    @pytest.mark.xfail(reason="no .ctypes attribute")
+    @pytest.mark.skip(reason="no .ctypes attribute")
     @pytest.mark.parametrize("dtype", [np.int32])
     def test__deepcopy__(self, dtype):
         # Force the entry of NULLs into array
@@ -1715,7 +1715,7 @@ def test_argsort(self):
                 assert_equal(a.copy().argsort(kind=kind), a, msg)
                 assert_equal(b.copy().argsort(kind=kind), b, msg)
 
-    @pytest.mark.xfail(reason='argsort complex')
+    @pytest.mark.skip(reason='argsort complex')
     def test_argsort_complex(self):
         a = np.arange(101, dtype=np.float32)
         b = np.flip(a)
@@ -2546,7 +2546,7 @@ def test_dot_out_mem_overlap(self):
             assert_raises(ValueError, np.dot, a, b, out=b[::2])
             assert_raises(ValueError, np.dot, a, b, out=b.T)
 
-    @pytest.mark.xfail(reason="TODO [::-1]")
+    @pytest.mark.xfail(reason="TODO: overlapping memor in matmul")
     def test_matmul_out(self):
         # overlapping memory
         a = np.arange(18).reshape(2, 3, 3)
@@ -2568,14 +2568,14 @@ def test_diagonal(self):
         assert_raises(np.AxisError, a.diagonal, axis1=0, axis2=5)
         assert_raises(np.AxisError, a.diagonal, axis1=5, axis2=0)
         assert_raises(np.AxisError, a.diagonal, axis1=5, axis2=5)
-        assert_raises(ValueError, a.diagonal, axis1=1, axis2=1)
+        assert_raises((ValueError, RuntimeError), a.diagonal, axis1=1, axis2=1)
 
         b = np.arange(8).reshape((2, 2, 2))
         assert_equal(b.diagonal(), [[0, 6], [1, 7]])
         assert_equal(b.diagonal(0), [[0, 6], [1, 7]])
         assert_equal(b.diagonal(1), [[2], [3]])
         assert_equal(b.diagonal(-1), [[4], [5]])
-        assert_raises(ValueError, b.diagonal, axis1=0, axis2=0)
+        assert_raises((ValueError, RuntimeError), b.diagonal, axis1=0, axis2=0)
         assert_equal(b.diagonal(0, 1, 2), [[0, 3], [4, 7]])
         assert_equal(b.diagonal(0, 0, 1), [[0, 6], [1, 7]])
         assert_equal(b.diagonal(offset=1, axis1=0, axis2=2), [[1], [3]])
@@ -2805,7 +2805,6 @@ def test_swapaxes(self):
                     if k == 1:
                         b = c
 
-    @pytest.mark.xfail(reason="TODO: ndarray.conjugate")
     def test_conjugate(self):
         a = np.array([1-1j, 1+1j, 23+23.0j])
         ac = a.conj()
@@ -2833,17 +2832,8 @@ def test_conjugate(self):
         assert_equal(ac, a.conjugate())
         assert_equal(ac, np.conjugate(a))
 
-        a = np.array([1-1j, 1+1j, 1, 2.0], object)
-        ac = a.conj()
-        assert_equal(ac, [k.conjugate() for k in a])
-        assert_equal(ac, a.conjugate())
-        assert_equal(ac, np.conjugate(a))
 
-        a = np.array([1-1j, 1, 2.0, 'f'], object)
-        assert_raises(TypeError, lambda: a.conj())
-        assert_raises(TypeError, lambda: a.conjugate())
-
-    @pytest.mark.xfail(reason="TODO: ndarray.conjugate")
+    @pytest.mark.xfail(reason="TODO: ndarray.conjugate with out")
     def test_conjugate_out(self):
         # Minimal test for the out argument being passed on correctly
         # NOTE: The ability to pass `out` is currently undocumented!
@@ -3754,7 +3744,7 @@ def test_ret_is_out(self, ndim, method):
         ret = arg_method(axis=0, out=out)
         assert ret is out
 
-    @pytest.mark.xfail(reason='FIXME: keepdims w/ positional args?')
+    @pytest.mark.xfail(reason='FIXME: out w/ positional args?')
     @pytest.mark.parametrize('arr_method, np_method',
         [('argmax', np.argmax),
          ('argmin', np.argmin)])
@@ -5438,7 +5428,7 @@ class TestDot:
     def setup_method(self):
         np.random.seed(128)
 
-        # Numpy guarantees the random stream, and we don't. So inline the
+        # Numpy and pytorch random streams differ, so inline the
         # values from numpy 1.24.1
         # self.A = np.random.rand(4, 2)
         self.A = np.array([[0.86663704, 0.26314485],
@@ -5626,7 +5616,7 @@ def test_dot_3args_errors(self):
         r = np.empty((1024, 32), dtype=int)
         assert_raises(ValueError, dot, f, v, r)
 
-    @pytest.mark.skip(reason="TODO order='F'")
+    @pytest.mark.xfail(reason="TODO order='F'")
     def test_dot_array_order(self):
         a = np.array([[1, 2], [3, 4]], order='C')
         b = np.array([[1, 2], [3, 4]], order='F')