Merge pull request #51 from Quansight-Labs/infer-full-dtype-followup

Make ones/zeros/empty/full dtype handling more uniform

Author: ev-br

torch_np/_dtypes.py

@@ -197,6 +197,7 @@ class bool_(generic):
 _typecodes = {st.typecode: st for cat in sctypes for st in sctypes[cat]}
 _torch_dtypes = {st.torch_dtype: st for cat in sctypes for st in sctypes[cat]}
 
+
 _aliases = {
     "u1": uint8,
     "i1": int8,
@@ -285,6 +286,11 @@ def name(self):
     def type(self):
         return self._scalar_type
 
+    @property
+    def kind(self):
+        # https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
+        return _torch_dtypes[self.torch_dtype].name[0]
+
     @property
     def typecode(self):
         return self._scalar_type.typecode

torch_np/_ndarray.py

@@ -493,13 +493,6 @@ def wrapped(x, *args, **kwds):
 ###### dtype routines
 
 
-def can_cast(from_, to, casting="safe"):
-    from_ = from_.dtype if isinstance(from_, ndarray) else _dtypes.dtype(from_)
-    to_ = to.dtype if isinstance(to, ndarray) else _dtypes.dtype(to)
-
-    return _dtypes_impl.can_cast_impl(from_.torch_dtype, to_.torch_dtype, casting)
-
-
 def _extract_dtype(entry):
     try:
         dty = _dtypes.dtype(entry)
@@ -508,6 +501,13 @@ def _extract_dtype(entry):
     return dty
 
 
+def can_cast(from_, to, casting="safe"):
+    from_ = _extract_dtype(from_)
+    to_ = _extract_dtype(to)
+
+    return _dtypes_impl.can_cast_impl(from_.torch_dtype, to_.torch_dtype, casting)
+
+
 def result_type(*arrays_and_dtypes):
     dtypes = []
 

torch_np/_wrapper.py

@@ -288,12 +288,15 @@ 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:
+        dtype = _dtypes_impl.default_float_dtype
+    result = torch.empty(shape, dtype=dtype)
+    return asarray(result)
 
 
 # NB: *_like function deliberately deviate from numpy: it has subok=True
@@ -310,17 +313,22 @@ def empty_like(prototype, dtype=None, order="K", subok=False, shape=None):
     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
-    if isinstance(fill_value, ndarray):
-        fill_value = fill_value.get()
+
+    fill_value = asarray(fill_value).get()
     if dtype is None:
-        torch_dtype = asarray(fill_value).get().dtype
-    else:
-        torch_dtype = _dtypes.torch_dtype_from(dtype)
-    return asarray(torch.full(shape, fill_value, dtype=torch_dtype))
+        dtype = fill_value.dtype
+
+    if not isinstance(shape, (tuple, list)):
+        shape = (shape,)
+
+    result = torch.full(shape, fill_value, dtype=dtype)
+
+    return asarray(result)
 
 
 @asarray_replacer()
@@ -335,12 +343,15 @@ 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:
+        dtype = _dtypes_impl.default_float_dtype
+    result = torch.ones(shape, dtype=dtype)
+    return asarray(result)
 
 
 @asarray_replacer()
@@ -362,7 +373,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()