MAINT: add a utility for to wrap tensor.to(dtype != tensor.dtype)

Author: ev-br

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

torch_np/_detail/implementations.py

@@ -33,8 +33,8 @@ def tensor_equiv(a1_t, a2_t):
 
 def isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
     dtype = _dtypes_impl.result_type_impl((a.dtype, b.dtype))
-    a = a.to(dtype)
-    b = b.to(dtype)
+    a = _util.cast_if_needed(a, dtype)
+    b = _util.cast_if_needed(b, dtype)
     result = torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
     return result
 
@@ -308,9 +308,7 @@ def corrcoef(xy_tensor, *, dtype=None):
         # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
         dtype = torch.float32
 
-    if dtype is not None:
-        xy_tensor = xy_tensor.to(dtype)
-
+    xy_tensor = _util.cast_if_needed(xy_tensor, dtype)
     result = torch.corrcoef(xy_tensor)
 
     if is_half:
@@ -336,8 +334,7 @@ def cov(
         # work around torch's "addmm_impl_cpu_" not implemented for 'Half'"
         dtype = torch.float32
 
-    if dtype is not None:
-        m_tensor = m_tensor.to(dtype)
+    m_tensor = _util.cast_if_needed(m_tensor, dtype)
 
     result = torch.cov(
         m_tensor, correction=ddof, aweights=aweights_tensor, fweights=fweights_tensor