MAINT: move shape/axes manipulations from ndarray

Author: ev-br

torch_np/_detail/implementations.py

@@ -581,21 +581,18 @@ def squeeze(tensor, axis=None):
     return result
 
 
-def reshape(tensor, *shape, order="C"):
+def reshape(tensor, shape, order="C"):
     if order != "C":
         raise NotImplementedError
-    newshape = shape[0] if len(shape) == 1 else shape
-    # convert any tnp.ndarray inputs into tensors before passing to torch.Tensor.reshape
-    t_newshape = _helpers.ndarrays_to_tensors(newshape)
     # if sh = (1, 2, 3), numpy allows both .reshape(sh) and .reshape(*sh)
-    result = tensor.reshape(t_newshape)
+    newshape = shape[0] if len(shape) == 1 else shape
+    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:
+def transpose(tensor, axes=None):
+    # numpy allows both .tranpose(sh) and .transpose(*sh)
+    if axes in [(), None, (None,)]:
         axes = tuple(range(tensor.ndim))[::-1]
     try:
         result = tensor.permute(axes)
@@ -604,6 +601,31 @@ def transpose(tensor, *axes):
     return result
 
 
+def ravel(tensor, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    result = tensor.ravel()
+    return result
+
+
+# leading underscore since arr.flatten exists but np.flatten does not
+def _flatten(tensor, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    # return a copy
+    result = tensor.ravel().clone()
+    return result
+
+
+# ### swap/move/roll axis ###
+
+def moveaxis(tensor, source, destination):
+    source = _util.normalize_axis_tuple(source, tensor.ndim, "source")
+    destination = _util.normalize_axis_tuple(destination, tensor.ndim, "destination")
+    result = torch.moveaxis(tensor, source, destination)
+    return result
+
+
 # ### Numeric ###
 
 

torch_np/_funcs.py

@@ -1,7 +1,7 @@
 import torch
 
 from . import _decorators, _helpers
-from ._detail import _util
+from ._detail import _util, _flips
 from ._detail import implementations as _impl
 
 
@@ -114,3 +114,60 @@ def searchsorted(a, v, side="left", sorter=None):
     a_t, v_t, sorter_t = _helpers.to_tensors_or_none(a, v, sorter)
     result = torch.searchsorted(a_t, v_t, side=side, sorter=sorter_t)
     return _helpers.array_from(result)
+
+
+# ### swap/move/roll axis ###
+
+
+def moveaxis(a, source, destination):
+    tensor, = _helpers.to_tensors(a)
+    result = _impl.moveaxis(tensor, source, destination)
+    return _helpers.array_from(result)
+
+
+def swapaxes(a, axis1, axis2):
+    tensor, = _helpers.to_tensors(a)
+    result = _flips.swapaxes(tensor, axis1, axis2)
+    return _helpers.array_from(result)
+
+
+def rollaxis(a, axis, start=0):
+    tensor, = _helpers.to_tensors(a)
+    result = _flips.rollaxis(a, axis, start)
+    return _helpers.array_from(result)
+
+
+# ### shape manipulations ###
+
+def squeeze(a, axis=None):
+    tensor, = _helpers.to_tensors(a)
+    result = _impl.squeeze(tensor, axis)
+    return _helpers.array_from(result, a)
+
+
+def reshape(a, newshape, order="C"):
+    tensor, = _helpers.to_tensors(a)    
+    result = _impl.reshape(tensor, newshape, order=order)
+    return _helpers.array_from(result, a)
+
+
+def transpose(a, axes=None):
+    tensor, = _helpers.to_tensors(a)    
+    result = _impl.transpose(tensor, axes)
+    return _helpers.array_from(result, a)
+
+
+def ravel(a, order="C"):
+    tensor, = _helpers.to_tensors(a)    
+    result = _impl.ravel(tensor)
+    return _helpers.array_from(result, a)  
+
+
+# leading underscore since arr.flatten exists but np.flatten does not
+def _flatten(a, order="C"):
+    tensor, = _helpers.to_tensors(a)    
+    result = _impl._flatten(tensor)
+    return _helpers.array_from(result, a)  
+
+
+

torch_np/_helpers.py

@@ -73,10 +73,10 @@ def result_or_out(result_tensor, out_array=None, promote_scalar=False):
         return asarray(result_tensor)
 
 
-def array_from(tensor):
+def array_from(tensor, base=None):
     from ._ndarray import ndarray
-
-    return ndarray._from_tensor_and_base(tensor, None)  # XXX: nuke .base
+    base = base if isinstance(base, ndarray) else None
+    return ndarray._from_tensor_and_base(tensor, base)  # XXX: nuke .base
 
 
 def tuple_arrays_from(result):

torch_np/_ndarray.py

@@ -348,31 +348,19 @@ def __irshift__(self, other):
 
     ### methods to match namespace functions
 
-    def squeeze(self, axis=None):
-        result = _impl.squeeze(self._tensor, axis)
-        return ndarray._from_tensor_and_base(result, self)
-
-    def reshape(self, *shape, order="C"):
-        result = _impl.reshape(self._tensor, *shape, order=order)
-        return ndarray._from_tensor_and_base(result, self)
+    squeeze = _funcs.squeeze
+    swapaxes = _funcs.swapaxes
 
     def transpose(self, *axes):
-        result = _impl.transpose(self._tensor, *axes)
-        return ndarray._from_tensor_and_base(result, self)
+        # np.transpose(arr, axis=None) but arr.transpose(*axes)
+        return _funcs.transpose(self, axes)
 
-    def swapaxes(self, axis1, axis2):
-        return asarray(_flips.swapaxes(self._tensor, axis1, axis2))
-
-    def ravel(self, order="C"):
-        if order != "C":
-            raise NotImplementedError
-        return ndarray._from_tensor_and_base(self._tensor.ravel(), self)
+    def reshape(self, *shape, order="C"):
+        # arr.reshape(shape) and arr.reshape(*shape)
+        return _funcs.reshape(self, shape, order=order)
 
-    def flatten(self, order="C"):
-        if order != "C":
-            raise NotImplementedError
-        result = self._tensor.flatten()
-        return asarray(result)
+    ravel = _funcs.ravel
+    flatten = _funcs._flatten
 
     nonzero = _funcs.nonzero
     clip = _funcs.clip

torch_np/_wrapper.py

@@ -457,26 +457,11 @@ def size(a, axis=None):
 ###### shape manipulations and indexing
 
 
-def transpose(a, axes=None):
-    arr = asarray(a)
-    return arr.transpose(axes)
-
-
-def reshape(a, newshape, order="C"):
-    arr = asarray(a)
-    return arr.reshape(*newshape, order=order)
-
-
 def ravel(a, order="C"):
     arr = asarray(a)
     return arr.ravel(order=order)
 
 
-def squeeze(a, axis=None):
-    arr = asarray(a)
-    return arr.squeeze(axis)
-
-
 def expand_dims(a, axis):
     a = asarray(a)
     shape = _util.expand_shape(a.shape, axis)
@@ -521,23 +506,6 @@ def broadcast_arrays(*args, subok=False):
     return tuple(asarray(_) for _ in res)
 
 
-@asarray_replacer()
-def moveaxis(a, source, destination):
-    source = _util.normalize_axis_tuple(source, a.ndim, "source")
-    destination = _util.normalize_axis_tuple(destination, a.ndim, "destination")
-    return asarray(torch.moveaxis(a, source, destination))
-
-
-def swapaxes(a, axis1, axis2):
-    arr = asarray(a)
-    return arr.swapaxes(axis1, axis2)
-
-
-@asarray_replacer()
-def rollaxis(a, axis, start=0):
-    return _flips.rollaxis(a, axis, start)
-
-
 def unravel_index(indices, shape, order="C"):
     # cf https://github.com/pytorch/pytorch/pull/66687
     # this version is from