BUG: random: size=() returns a 0D array, not scalar

Author: ev-br

torch_np/random.py

@@ -29,8 +29,8 @@
 ]
 
 
-def array_or_scalar(values, py_type=float):
-    if values.numel() == 1:
+def array_or_scalar(values, py_type=float, size=None):
+    if size is None:
         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)
+    return array_or_scalar(values, size=size)
 
 
 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)
+    return array_or_scalar(values, size=size)
 
 
 def randn(*size):
     values = torch.randn(size, dtype=_default_dtype)
-    return array_or_scalar(values)
+    return array_or_scalar(values, size=size)
 
 
 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)
+    return array_or_scalar(values, size=size)
 
 
 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)
+    return array_or_scalar(values, int, size=size)
 
 
 def choice(a, size=None, replace=True, p=None):

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -3627,7 +3627,6 @@ def test_assign_mask2(self):
         assert_array_equal(x, np.array([[1, 10, 3, 4], [5, 6, 7, 8]]))
 
 
-@pytest.mark.xfail(reason='TODO')
 class TestArgmaxArgminCommon:
 
     sizes = [(), (3,), (3, 2), (2, 3),