Merge pull request #69 from Quansight-Labs/free_funcs

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/_detail/implementations.py

@@ -1,6 +1,5 @@
 import torch
 
-from .. import _helpers
 from . import _dtypes_impl, _util
 
 # ### equality, equivalence, allclose ###
@@ -581,21 +580,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 +600,32 @@ 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.flatten()
+    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 ###
 
 
@@ -622,6 +644,14 @@ def round(tensor, decimals=0):
     return result
 
 
+def imag(tensor):
+    if tensor.is_complex():
+        result = tensor.imag
+    else:
+        result = torch.zeros_like(tensor)
+    return result
+
+
 # ### put/take along axis ###
 
 

torch_np/_funcs.py

@@ -0,0 +1,196 @@
+import torch
+
+from . import _decorators, _helpers
+from ._detail import _flips, _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)
+
+
+# ### 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)
+
+
+# ### Type/shape etc queries ###
+
+
+def real(a):
+    (tensor,) = _helpers.to_tensors(a)
+    result = torch.real(tensor)
+    return _helpers.array_from(result)
+
+
+def imag(a):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.imag(tensor)
+    return _helpers.array_from(result)
+
+
+def round_(a, decimals=0, out=None):
+    (tensor,) = _helpers.to_tensors(a)
+    result = _impl.round(tensor, decimals)
+    return _helpers.result_or_out(result, out)
+
+
+around = round_
+round = round_

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,22 @@ def result_or_out(result_tensor, out_array=None, promote_scalar=False):
         return asarray(result_tensor)
 
 
+def array_from(tensor, base=None):
+    from ._ndarray import ndarray
+
+    base = base if isinstance(base, ndarray) else None
+    return ndarray._from_tensor_and_base(tensor, base)  # 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
@@ -105,11 +124,3 @@ def to_tensors_or_none(*inputs):
     from ._ndarray import asarray, ndarray
 
     return tuple(None if value is None else asarray(value).get() for value in inputs)
-
-
-def _outer(x, y):
-    from ._ndarray import asarray
-
-    x_tensor, y_tensor = to_tensors(x, y)
-    result = torch.outer(x_tensor, y_tensor)
-    return asarray(result)