MAINT: merge torch implementations of _flips.py to _flips

Author: ev-br

torch_np/_detail/__init__.py

@@ -3,4 +3,3 @@
 
 # leading underscore (ndarray.flatten yes, np.flatten no)
 from .implementations import *
-from .implementations import _flatten

torch_np/_detail/_flips.py

@@ -26,36 +26,3 @@ def fliplr(m_tensor):
 def rot90(m_tensor, k=1, axes=(0, 1)):
     axes = _util.normalize_axis_tuple(axes, m_tensor.ndim)
     return torch.rot90(m_tensor, k, axes)
-
-
-def swapaxes(tensor, axis1, axis2):
-    axis1 = _util.normalize_axis_index(axis1, tensor.ndim)
-    axis2 = _util.normalize_axis_index(axis2, tensor.ndim)
-    return torch.swapaxes(tensor, axis1, axis2)
-
-
-# Straight vendor from:
-# https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1259
-#
-# Also note this function in NumPy is mostly retained for backwards compat
-# (https://stackoverflow.com/questions/29891583/reason-why-numpy-rollaxis-is-so-confusing)
-# so let's not touch it unless hard pressed.
-def rollaxis(tensor, axis, start=0):
-    n = tensor.ndim
-    axis = _util.normalize_axis_index(axis, n)
-    if start < 0:
-        start += n
-    msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
-    if not (0 <= start < n + 1):
-        raise _util.AxisError(msg % ("start", -n, "start", n + 1, start))
-    if axis < start:
-        # it's been removed
-        start -= 1
-    if axis == start:
-        # numpy returns a view, here we try returning the tensor itself
-        # return tensor[...]
-        return tensor
-    axes = list(range(0, n))
-    axes.remove(axis)
-    axes.insert(start, axis)
-    return tensor.view(axes)

torch_np/_detail/implementations.py

@@ -129,62 +129,6 @@ def triu_indices(n, k=0, m=None):
     return result
 
 
-def diag_indices(n, ndim=2):
-    idx = torch.arange(n)
-    return (idx,) * ndim
-
-
-def diag_indices_from(tensor):
-    if not tensor.ndim >= 2:
-        raise ValueError("input array must be at least 2-d")
-    # For more than d=2, the strided formula is only valid for arrays with
-    # all dimensions equal, so we check first.
-    s = tensor.shape
-    if s[1:] != s[:-1]:
-        raise ValueError("All dimensions of input must be of equal length")
-    return diag_indices(s[0], tensor.ndim)
-
-
-def fill_diagonal(tensor, t_val, wrap):
-    # torch.Tensor.fill_diagonal_ only accepts scalars. Thus vendor the numpy source,
-    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/lib/index_tricks.py#L786-L917
-
-    if tensor.ndim < 2:
-        raise ValueError("array must be at least 2-d")
-    end = None
-    if tensor.ndim == 2:
-        # Explicit, fast formula for the common case.  For 2-d arrays, we
-        # accept rectangular ones.
-        step = tensor.shape[1] + 1
-        # This is needed to don't have tall matrix have the diagonal wrap.
-        if not wrap:
-            end = tensor.shape[1] * tensor.shape[1]
-    else:
-        # For more than d=2, the strided formula is only valid for arrays with
-        # all dimensions equal, so we check first.
-        s = tensor.shape
-        if s[1:] != s[:-1]:
-            raise ValueError("All dimensions of input must be of equal length")
-        sz = torch.as_tensor(tensor.shape[:-1])
-        step = 1 + (torch.cumprod(sz, 0)).sum()
-
-    # Write the value out into the diagonal.
-    tensor.ravel()[:end:step] = t_val
-    return tensor
-
-
-def trace(tensor, offset=0, axis1=0, axis2=1, dtype=None, out=None):
-    result = torch.diagonal(tensor, offset, dim1=axis1, dim2=axis2).sum(-1, dtype=dtype)
-    return result
-
-
-def diagonal(tensor, offset=0, axis1=0, axis2=1):
-    axis1 = _util.normalize_axis_index(axis1, tensor.ndim)
-    axis2 = _util.normalize_axis_index(axis2, tensor.ndim)
-    result = torch.diagonal(tensor, offset, axis1, axis2)
-    return result
-
-
 # ### splits ###
 
 
@@ -263,14 +207,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 None and t_max is None:
-        raise ValueError("One of max or min must be given")
-
-    result = tensor.clamp(t_min, t_max)
-    return result
-
-
 def diff(a_tensor, n=1, axis=-1, prepend_tensor=None, append_tensor=None):
     axis = _util.normalize_axis_index(axis, a_tensor.ndim)
 
@@ -364,14 +300,6 @@ def vstack(tensors, *, dtype=None, casting="same_kind"):
     return result
 
 
-def tile(tensor, reps):
-    if isinstance(reps, int):
-        reps = (reps,)
-
-    result = torch.tile(tensor, reps)
-    return result
-
-
 # #### cov & corrcoef
 
 
@@ -549,14 +477,6 @@ def arange(start=None, stop=None, step=1, dtype=None):
 # ### empty/full et al ###
 
 
-def eye(N, M=None, k=0, dtype=float):
-    if M is None:
-        M = N
-    z = torch.zeros(N, M, dtype=dtype)
-    z.diagonal(k).fill_(1)
-    return z
-
-
 def zeros(shape, dtype=None, order="C"):
     if order != "C":
         raise NotImplementedError
@@ -637,102 +557,12 @@ def full(shape, fill_value, dtype=None, order="C"):
 # ### 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:
-        if isinstance(axis, tuple):
-            result = tensor
-            for ax in axis:
-                result = result.squeeze(ax)
-        else:
-            result = tensor.squeeze(axis)
-    return result
-
-
-def reshape(tensor, shape, order="C"):
-    if order != "C":
-        raise NotImplementedError
-    # if sh = (1, 2, 3), numpy allows both .reshape(sh) and .reshape(*sh)
-    newshape = shape[0] if len(shape) == 1 else shape
-    result = tensor.reshape(newshape)
-    return result
-
-
-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)
-    except RuntimeError:
-        raise ValueError("axes don't match array")
-    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 ###
 
 
-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
-
-
-def imag(tensor):
-    if tensor.is_complex():
-        result = tensor.imag
-    else:
-        result = torch.zeros_like(tensor)
-    return result
-
-
 # ### put/take along axis ###
 
 
@@ -753,36 +583,6 @@ def put_along_dim(tensor, t_indices, t_values, axis):
     return result
 
 
-# ### sort and partition ###
-
-
-def _sort_helper(tensor, axis, kind, order):
-    if order is not None:
-        # only relevant for structured dtypes; not supported
-        raise NotImplementedError(
-            "'order' keyword is only relevant for structured dtypes"
-        )
-
-    (tensor,), axis = _util.axis_none_ravel(tensor, axis=axis)
-    axis = _util.normalize_axis_index(axis, tensor.ndim)
-
-    stable = kind == "stable"
-
-    return tensor, axis, stable
-
-
-def sort(tensor, axis=-1, kind=None, order=None):
-    tensor, axis, stable = _sort_helper(tensor, axis, kind, order)
-    result = torch.sort(tensor, dim=axis, stable=stable)
-    return result.values
-
-
-def argsort(tensor, axis=-1, kind=None, order=None):
-    tensor, axis, stable = _sort_helper(tensor, axis, kind, order)
-    result = torch.argsort(tensor, dim=axis, stable=stable)
-    return result
-
-
 # ### logic and selection ###
 
 
@@ -831,56 +631,6 @@ def inner(t_a, t_b):
     return result
 
 
-def vdot(t_a, t_b, /):
-    # 1. torch only accepts 1D arrays, numpy ravels
-    # 2. torch requires matching dtype, while numpy casts (?)
-    t_a, t_b = torch.atleast_1d(t_a, t_b)
-    if t_a.ndim > 1:
-        t_a = t_a.ravel()
-    if t_b.ndim > 1:
-        t_b = t_b.ravel()
-
-    dtype = _dtypes_impl.result_type_impl((t_a.dtype, t_b.dtype))
-    is_half = dtype == torch.float16
-    is_bool = dtype == torch.bool
-
-    # work around torch's "dot" not implemented for 'Half', 'Bool'
-    if is_half:
-        dtype = torch.float32
-    if is_bool:
-        dtype = torch.uint8
-
-    t_a = _util.cast_if_needed(t_a, dtype)
-    t_b = _util.cast_if_needed(t_b, dtype)
-
-    result = torch.vdot(t_a, t_b)
-
-    if is_half:
-        result = result.to(torch.float16)
-    if is_bool:
-        result = result.to(torch.bool)
-
-    return result
-
-
-def dot(t_a, t_b):
-    dtype = _dtypes_impl.result_type_impl((t_a.dtype, t_b.dtype))
-    t_a = _util.cast_if_needed(t_a, dtype)
-    t_b = _util.cast_if_needed(t_b, dtype)
-
-    if t_a.ndim == 0 or t_b.ndim == 0:
-        result = t_a * t_b
-    elif t_a.ndim == 1 and t_b.ndim == 1:
-        result = torch.dot(t_a, t_b)
-    elif t_a.ndim == 1:
-        result = torch.mv(t_b.T, t_a).T
-    elif t_b.ndim == 1:
-        result = torch.mv(t_a, t_b)
-    else:
-        result = torch.matmul(t_a, t_b)
-    return result
-
-
 # ### unique et al ###