Merge pull request #52 from Quansight-Labs/move_to_impl

make better split between wrappers and torch implementations

Author: ev-br

torch_np/_decorators.py

@@ -15,7 +15,7 @@ def wrapped(*args, dtype=None, **kwds):
         torch_dtype = None
         if dtype is not None:
             dtype = _dtypes.dtype(dtype)
-            torch_dtype = dtype._scalar_type.torch_dtype
+            torch_dtype = dtype.torch_dtype
         return func(*args, dtype=torch_dtype, **kwds)
 
     return wrapped

torch_np/_detail/_reductions.py

@@ -146,9 +146,7 @@ def std(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
         raise NotImplementedError
 
     dtype = _atleast_float(dtype, tensor.dtype)
-
-    if dtype is not None:
-        tensor = tensor.to(dtype)
+    tensor = _util.cast_if_needed(tensor, dtype)
     result = tensor.std(dim=axis, correction=ddof)
 
     return result
@@ -159,9 +157,7 @@ def var(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
         raise NotImplementedError
 
     dtype = _atleast_float(dtype, tensor.dtype)
-
-    if dtype is not None:
-        tensor = tensor.to(dtype)
+    tensor = _util.cast_if_needed(tensor, dtype)
     result = tensor.var(dim=axis, correction=ddof)
 
     return result
@@ -204,10 +200,9 @@ def average(a_tensor, axis, w_tensor):
         a_tensor = a_tensor.to(result_dtype)
 
     result_dtype = _dtypes_impl.result_type_impl([a_tensor.dtype, w_tensor.dtype])
-    if a_tensor.dtype != result_dtype:
-        a_tensor = a_tensor.to(result_dtype)
-    if w_tensor.dtype != result_dtype:
-        w_tensor = w_tensor.to(result_dtype)
+
+    a_tensor = _util.cast_if_needed(a_tensor, result_dtype)
+    w_tensor = _util.cast_if_needed(w_tensor, result_dtype)
 
     # axis
     if axis is None:
@@ -258,7 +253,7 @@ def quantile(a_tensor, q_tensor, axis, method):
         axis = _util.normalize_axis_tuple(axis, a_tensor.ndim)
     axis = _util.allow_only_single_axis(axis)
 
-    q_tensor = q_tensor.to(a_tensor.dtype)
+    q_tensor = _util.cast_if_needed(q_tensor, a_tensor.dtype)
 
     (a_tensor, q_tensor), axis = _util.axis_none_ravel(a_tensor, q_tensor, axis=axis)
 

torch_np/_detail/_util.py

@@ -34,6 +34,13 @@ class UFuncTypeError(TypeError, RuntimeError):
     pass
 
 
+def cast_if_needed(tensor, dtype):
+    # NB: no casting if dtype=None
+    if tensor.dtype != dtype:
+        tensor = tensor.to(dtype)
+    return tensor
+
+
 # a replica of the version in ./numpy/numpy/core/src/multiarray/common.h
 def normalize_axis_index(ax, ndim, argname=None):
     if not (-ndim <= ax < ndim):
@@ -156,10 +163,7 @@ def cast_dont_broadcast(tensors, target_dtype, casting):
                 f"Cannot cast array data from {tensor.dtype} to"
                 f" {target_dtype} according to the rule '{casting}'"
             )
-
-        # cast if needed
-        if tensor.dtype != target_dtype:
-            tensor = tensor.to(target_dtype)
+        tensor = cast_if_needed(tensor, target_dtype)
         cast_tensors.append(tensor)
 
     return tuple(cast_tensors)
@@ -200,8 +204,7 @@ def cast_and_broadcast(tensors, out_param, casting):
             )
 
         # cast arr if needed
-        if tensor.dtype != target_dtype:
-            tensor = tensor.to(target_dtype)
+        tensor = cast_if_needed(tensor, target_dtype)
 
         # `out` broadcasts `tensor`
         if tensor.shape != target_shape:
@@ -285,8 +288,7 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
         tensor = torch.as_tensor(obj, dtype=torch_dtype)
 
     # type cast if requested
-    if dtype is not None:
-        tensor = tensor.to(dtype)
+    tensor = cast_if_needed(tensor, dtype)
 
     # adjust ndim if needed
     ndim_extra = ndmin - tensor.ndim