MAINT: move tensor manipulations from _ndarray to _impl

Author: ev-br

torch_np/_detail/implementations.py

@@ -216,7 +216,6 @@ def dsplit(tensor, indices_or_sections):
     return split_helper(tensor, indices_or_sections, 2, strict=True)
 
 
-
 def clip(tensor, t_min, t_max):
     if t_min is not None:
         t_min = torch.broadcast_to(t_min, tensor.shape)
@@ -488,3 +487,64 @@ def full(shape, fill_value, dtype=None):
         shape = (shape,)
     result = torch.full(shape, fill_value, dtype=dtype)
     return result
+
+
+# ### shape manipulations ###
+
+
+def roll(tensor, shift, axis=None):
+    if axis is not None:
+        axis = _util.normalize_axis_tuple(axis, tensor.ndim, allow_duplicate=True)
+        if not isinstance(shift, tuple):
+            shift = (shift,) * len(axis)
+    result = tensor.roll(shift, axis)
+    return result
+
+
+def squeeze(tensor, axis=None):
+    if axis == ():
+        result = tensor
+    elif axis is None:
+        result = tensor.squeeze()
+    else:
+        result = tensor.squeeze(axis)
+    return result
+
+
+def reshape(tensor, *shape, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    newshape = shape[0] if len(shape) == 1 else shape
+    # if sh = (1, 2, 3), numpy allows both .reshape(sh) and .reshape(*sh)
+    result = tensor.reshape(newshape)
+    return result
+
+
+def transpose(tensor, *axes):
+    # numpy allows both .reshape(sh) and .reshape(*sh)
+    axes = axes[0] if len(axes) == 1 else axes
+    if axes == () or axes is None:
+        axes = tuple(range(tensor.ndim))[::-1]
+    try:
+        result = tensor.permute(axes)
+    except RuntimeError:
+        raise ValueError("axes don't match array")
+    return result
+
+
+# ### Numeric ###
+
+
+def round(tensor, decimals=0):
+    if tensor.is_floating_point():
+        result = torch.round(tensor, decimals=decimals)
+    elif tensor.is_complex():
+        # RuntimeError: "round_cpu" not implemented for 'ComplexFloat'
+        result = (
+            torch.round(tensor.real, decimals=decimals)
+            + torch.round(tensor.imag, decimals=decimals) * 1j
+        )
+    else:
+        # RuntimeError: "round_cpu" not implemented for 'int'
+        result = tensor
+    return result

torch_np/_ndarray.py

@@ -139,11 +139,7 @@ def imag(self, value):
         self._tensor.imag = asarray(value).get()
 
     def round(self, decimals=0, out=None):
-        tensor = self._tensor
-        if torch.is_floating_point(tensor):
-            result = torch.round(tensor, decimals=decimals)
-        else:
-            result = tensor
+        result = _impl.round(self._tensor, decimals)
         return _helpers.result_or_out(result, out)
 
     # ctors
@@ -328,32 +324,16 @@ def __irshift__(self, other):
     ### methods to match namespace functions
 
     def squeeze(self, axis=None):
-        if axis == ():
-            tensor = self._tensor
-        elif axis is None:
-            tensor = self._tensor.squeeze()
-        else:
-            tensor = self._tensor.squeeze(axis)
-        return ndarray._from_tensor_and_base(tensor, self)
+        result = _impl.squeeze(self._tensor, axis)
+        return ndarray._from_tensor_and_base(result, self)
 
     def reshape(self, *shape, order="C"):
-        newshape = shape[0] if len(shape) == 1 else shape
-        # if sh = (1, 2, 3), numpy allows both .reshape(sh) and .reshape(*sh)
-        if order != "C":
-            raise NotImplementedError
-        tensor = self._tensor.reshape(newshape)
-        return ndarray._from_tensor_and_base(tensor, self)
+        result = _impl.reshape(self._tensor, *shape, order=order)
+        return ndarray._from_tensor_and_base(result, self)
 
     def transpose(self, *axes):
-        # numpy allows both .reshape(sh) and .reshape(*sh)
-        axes = axes[0] if len(axes) == 1 else axes
-        if axes == () or axes is None:
-            axes = tuple(range(self.ndim))[::-1]
-        try:
-            tensor = self._tensor.permute(axes)
-        except RuntimeError:
-            raise ValueError("axes don't match array")
-        return ndarray._from_tensor_and_base(tensor, self)
+        result = _impl.transpose(self._tensor, *axes)
+        return ndarray._from_tensor_and_base(result, self)
 
     def swapaxes(self, axis1, axis2):
         return asarray(_flips.swapaxes(self._tensor, axis1, axis2))

torch_np/_wrapper.py

@@ -595,9 +595,9 @@ def flatnonzero(a):
 
 
 def argwhere(a):
-    arr = asarray(a)
-    tensor = arr.get()
-    return asarray(torch.argwhere(tensor))
+    tensor = asarray(a).get()
+    result = torch.argwhere(tensor)
+    return asarray(result)
 
 
 from ._decorators import emulate_out_arg
@@ -606,13 +606,10 @@ def argwhere(a):
 count_nonzero = emulate_out_arg(axis_keepdims_wrapper(_reductions.count_nonzero))
 
 
-@asarray_replacer()
 def roll(a, shift, axis=None):
-    if axis is not None:
-        axis = _util.normalize_axis_tuple(axis, a.ndim, allow_duplicate=True)
-        if not isinstance(shift, tuple):
-            shift = (shift,) * len(axis)
-    return a.roll(shift, axis)
+    tensor = asarray(a).get()
+    result = _impl.roll(tensor, shift, axis)
+    return asarray(result)
 
 
 def round_(a, decimals=0, out=None):