MAINT: empty/full/ones/zeros: dtype handling

Author: ev-br

torch_np/_wrapper.py

@@ -282,12 +282,20 @@ def arange(start=None, stop=None, step=1, dtype=None, *, like=None):
         raise ValueError("Maximum allowed size exceeded")
 
 
+@_decorators.dtype_to_torch
 def empty(shape, dtype=float, order="C", *, like=None):
     _util.subok_not_ok(like)
     if order != "C":
         raise NotImplementedError
-    torch_dtype = _dtypes.torch_dtype_from(dtype)
-    return asarray(torch.empty(shape, dtype=torch_dtype))
+
+    if dtype is None:
+        from ._detail._scalar_types import default_float_type
+
+        dtype = default_float_type.torch_dtype
+
+    result = torch.empty(shape, dtype=dtype)
+
+    return asarray(result)
 
 
 # NB: *_like function deliberately deviate from numpy: it has subok=True
@@ -303,15 +311,22 @@ def empty_like(prototype, dtype=None, order="K", subok=False, shape=None):
         result = result.reshape(shape)
     return result
 
+
 @_decorators.dtype_to_torch
 def full(shape, fill_value, dtype=None, order="C", *, like=None):
     _util.subok_not_ok(like)
     if order != "C":
         raise NotImplementedError
+
     fill_value = asarray(fill_value).get()
     if dtype is None:
-        dtype =  fill_value.dtype
+        dtype = fill_value.dtype
+
+    if not isinstance(shape, (tuple, list)):
+        shape = (shape,)
+
     result = torch.full(shape, fill_value, dtype=dtype)
+
     return asarray(result)
 
 
@@ -327,12 +342,19 @@ def full_like(a, fill_value, dtype=None, order="K", subok=False, shape=None):
     return result
 
 
+@_decorators.dtype_to_torch
 def ones(shape, dtype=None, order="C", *, like=None):
     _util.subok_not_ok(like)
     if order != "C":
         raise NotImplementedError
-    torch_dtype = _dtypes.torch_dtype_from(dtype)
-    return asarray(torch.ones(shape, dtype=torch_dtype))
+    if dtype is None:
+        from ._detail._scalar_types import default_float_type
+
+        dtype = default_float_type.torch_dtype
+
+    result = torch.ones(shape, dtype=dtype)
+
+    return asarray(result)
 
 
 @asarray_replacer()
@@ -354,7 +376,8 @@ def zeros(shape, dtype=None, order="C", *, like=None):
         raise NotImplementedError
     if dtype is None:
         dtype = _dtypes_impl.default_float_dtype
-    return asarray(torch.zeros(shape, dtype=dtype))
+    result = torch.zeros(shape, dtype=dtype)
+    return asarray(result)
 
 
 @asarray_replacer()

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -2666,12 +2666,19 @@ def test_mode(self):
             np.convolve(d, k, mode=None)
 
 
+class TestDtypePositional:
+
+    @pytest.mark.xfail(reason='TODO: restore dtypes as positional args')
+    def test_dtype_positional(self):
+        np.empty((2,), bool)
+
+
 class TestArgwhere:
 
     @pytest.mark.parametrize('nd', [0, 1, 2])
     def test_nd(self, nd):
         # get an nd array with multiple elements in every dimension
-        x = np.empty((2,)*nd, bool)
+        x = np.empty((2,)*nd, dtype=bool)
 
         # none
         x[...] = False

torch_np/tests/numpy_tests/lib/test_shape_base_.py

@@ -719,7 +719,7 @@ def test_kroncompare(self):
         for s in shape:
             b = randint(0, 10, size=s)
             for r in reps:
-                a = np.ones(r, b.dtype)
+                a = np.ones(r, dtype=b.dtype)   # TODO: restore dtype positional arg
                 large = tile(b, r)
                 klarge = kron(a, b)
                 assert_equal(large, klarge)