MAINT: move ndarray methods to free functions

Author: ev-br

torch_np/__init__.py

@@ -4,6 +4,7 @@
 from ._detail._index_tricks import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
+from ._funcs import *
 from ._getlimits import errstate, finfo, iinfo
 from ._ndarray import array, asarray, can_cast, ndarray, newaxis, result_type
 from ._unary_ufuncs import *

torch_np/_funcs.py

@@ -0,0 +1,116 @@
+import torch
+
+from . import _decorators, _helpers
+from ._detail import _util
+from ._detail import implementations as _impl
+
+
+def nonzero(a):
+    (tensor,) = _helpers.to_tensors(a)
+    result = tensor.nonzero(as_tuple=True)
+    return _helpers.tuple_arrays_from(result)
+
+
+def argwhere(a):
+    (tensor,) = _helpers.to_tensors(a)
+    result = torch.argwhere(tensor)
+    return _helpers.array_from(result)
+
+
+def clip(a, min=None, max=None, out=None):
+    # np.clip requires both a_min and a_max not None, while ndarray.clip allows
+    # one of them to be None. Follow the more lax version.
+    # Also min/max as arg names: follow numpy naming.
+    tensor, t_min, t_max = _helpers.to_tensors_or_none(a, min, max)
+    result = _impl.clip(tensor, t_min, t_max)
+    return _helpers.result_or_out(result, out)
+
+
+def repeat(a, repeats, axis=None):
+    tensor, t_repeats = _helpers.to_tensors(a, repeats)  # XXX: scalar repeats
+    result = torch.repeat_interleave(tensor, t_repeats, axis)
+    return _helpers.array_from(result)
+
+
+# ### diag et al ###
+
+
+def diagonal(a, offset=0, axis1=0, axis2=1):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.diagonal(tensor, offset, axis1, axis2)
+    return _helpers.array_from(result)
+
+
+@_decorators.dtype_to_torch
+def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.trace(tensor, offset, axis1, axis2, dtype)
+    return _helpers.result_or_out(result, out)
+
+
+@_decorators.dtype_to_torch
+def eye(N, M=None, k=0, dtype=float, order="C", *, like=None):
+    _util.subok_not_ok(like)
+    if order != "C":
+        raise NotImplementedError
+    result = _impl.eye(N, M, k, dtype)
+    return _helpers.array_from(result)
+
+
+@_decorators.dtype_to_torch
+def identity(n, dtype=None, *, like=None):
+    _util.subok_not_ok(like)
+    result = torch.eye(n, dtype=dtype)
+    return _helpers.array_from(result)
+
+
+def diag(v, k=0):
+    (tensor,) = _helpers.to_tensors(v)
+    result = torch.diag(tensor, k)
+    return _helpers.array_from(result)
+
+
+def diagflat(v, k=0):
+    (tensor,) = _helpers.to_tensors(v)
+    result = torch.diagflat(tensor, k)
+    return _helpers.array_from(result)
+
+
+def diag_indices(n, ndim=2):
+    result = _impl.diag_indices(n, ndim)
+    return _helpers.tuple_arrays_from(result)
+
+
+def diag_indices_from(arr):
+    (tensor,) = _helpers.to_tensors(arr)
+    result = _impl.diag_indices_from(tensor)
+    return _helpers.tuple_arrays_from(result)
+
+
+def fill_diagonal(a, val, wrap=False):
+    tensor, t_val = _helpers.to_tensors(a, val)
+    result = _impl.fill_diagonal(tensor, t_val, wrap)
+    return _helpers.array_from(result)
+
+
+# ### sorting ###
+
+# ### sort and partition ###
+
+
+def sort(a, axis=-1, kind=None, order=None):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.sort(tensor, axis, kind, order)
+    return _helpers.array_from(result)
+
+
+def argsort(a, axis=-1, kind=None, order=None):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.argsort(tensor, axis, kind, order)
+    return _helpers.array_from(result)
+
+
+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)

torch_np/_helpers.py

@@ -41,6 +41,9 @@ def cast_and_broadcast(tensors, out, casting):
     return tuple(tensors)
 
 
+# ### Return helpers: wrap a single tensor, a tuple of tensors, out= etc ###
+
+
 def result_or_out(result_tensor, out_array=None, promote_scalar=False):
     """A helper for returns with out= argument.
 
@@ -70,6 +73,21 @@ def result_or_out(result_tensor, out_array=None, promote_scalar=False):
         return asarray(result_tensor)
 
 
+def array_from(tensor):
+    from ._ndarray import ndarray
+
+    return ndarray._from_tensor_and_base(tensor, None)  # XXX: nuke .base
+
+
+def tuple_arrays_from(result):
+    from ._ndarray import asarray
+
+    return tuple(asarray(x) for x in result)
+
+
+# ### Various ways of converting array-likes to tensors ###
+
+
 def ndarrays_to_tensors(*inputs):
     """Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
     from ._ndarray import asarray, ndarray

torch_np/_ndarray.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from . import _binary_ufuncs, _dtypes, _helpers, _unary_ufuncs
+from . import _binary_ufuncs, _dtypes, _funcs, _helpers, _unary_ufuncs
 from ._decorators import (
     NoValue,
     axis_keepdims_wrapper,
@@ -374,20 +374,24 @@ def flatten(self, order="C"):
         result = self._tensor.flatten()
         return asarray(result)
 
-    def nonzero(self):
-        tensor = self._tensor
-        return tuple(asarray(_) for _ in tensor.nonzero(as_tuple=True))
+    nonzero = _funcs.nonzero
+    clip = _funcs.clip
+    repeat = _funcs.repeat
 
-    def clip(self, min=None, max=None, out=None):
-        tensor, t_min, t_max = _helpers.to_tensors_or_none(self, min, max)
-        result = _impl.clip(tensor, t_min, t_max)
-        return _helpers.result_or_out(result, out)
+    diagonal = _funcs.diagonal
+    trace = _funcs.trace
 
-    def repeat(self, repeats, axis=None):
-        t_repeats = asarray(repeats).get()  # XXX: scalar repeats
-        tensor = self._tensor
-        result = torch.repeat_interleave(tensor, t_repeats, axis)
-        return asarray(result)
+    ### sorting ###
+
+    def sort(self, axis=-1, kind=None, order=None):
+        # ndarray.sort works in-place
+        result = _impl.sort(self._tensor, axis, kind, order)
+        self._tensor = result
+
+    argsort = _funcs.argsort
+    searchsorted = _funcs.searchsorted
+
+    ### reductions ###
 
     argmin = emulate_out_arg(axis_keepdims_wrapper(_reductions.argmin))
     argmax = emulate_out_arg(axis_keepdims_wrapper(_reductions.argmax))
@@ -411,16 +415,6 @@ def repeat(self, repeats, axis=None):
         axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumsum))
     )
 
-    def diagonal(self, offset=0, axis1=0, axis2=1):
-        result = _impl.diagonal(self._tensor, offset, axis1, axis2)
-        return asarray(result)
-
-    @dtype_to_torch
-    def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
-        tensor = self._tensor
-        result = _impl.trace(tensor, offset, axis1, axis2, dtype)
-        return _helpers.result_or_out(result, out)
-
     ### indexing ###
     @staticmethod
     def _upcast_int_indices(index):
@@ -442,22 +436,6 @@ def __setitem__(self, index, value):
         value = _helpers.ndarrays_to_tensors(value)
         return self._tensor.__setitem__(index, value)
 
-    ### sorting ###
-
-    def sort(self, axis=-1, kind=None, order=None):
-        # ndarray.sort works in-place
-        result = _impl.sort(self._tensor, axis, kind, order)
-        self._tensor = result
-
-    def argsort(self, axis=-1, kind=None, order=None):
-        result = _impl.argsort(self._tensor, axis, kind, order)
-        return asarray(result)
-
-    def searchsorted(self, v, side="left", sorter=None):
-        v_t, sorter_t = _helpers.to_tensors_or_none(v, sorter)
-        result = torch.searchsorted(self._tensor, v_t, side=side, sorter=sorter_t)
-        return asarray(result)
-
 
 # This is the ideally the only place which talks to ndarray directly.
 # The rest goes through asarray (preferred) or array.