MAINT: split remaining functions into normalizations and implementations

Author: ev-br

torch_np/_detail/_reductions.py

@@ -371,7 +371,24 @@ def average_weights(a_tensor, axis, w_tensor, keepdims=False):
     return result, denominator
 
 
-def quantile(a_tensor, q_tensor, axis, method, keepdims=False):
+def quantile(
+    a_tensor,
+    q_tensor,
+    axis,
+    overwrite_input,
+    method,
+    keepdims=False,
+    interpolation=None,
+):
+    if overwrite_input:
+        # raise NotImplementedError("overwrite_input in quantile not implemented.")
+        # NumPy documents that `overwrite_input` MAY modify inputs:
+        # https://numpy.org/doc/stable/reference/generated/numpy.percentile.html#numpy-percentile
+        # Here we choose to work out-of-place because why not.
+        pass
+
+    if interpolation is not None:
+        raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
 
     if (0 > q_tensor).any() or (q_tensor > 1).any():
         raise ValueError("Quantiles must be in range [0, 1], got %s" % q_tensor)
@@ -404,3 +421,23 @@ def quantile(a_tensor, q_tensor, axis, method, keepdims=False):
     if keepdims:
         result = _util.apply_keepdims(result, ax, ndim)
     return result
+
+
+def percentile(
+    a_tensor,
+    q_tensor,
+    axis,
+    overwrite_input,
+    method,
+    keepdims=False,
+    interpolation=None,
+):
+    return quantile(
+        a_tensor,
+        q_tensor / 100.0,
+        axis=axis,
+        overwrite_input=overwrite_input,
+        method=method,
+        keepdims=keepdims,
+        interpolation=interpolation,
+    )

torch_np/_detail/implementations.py

@@ -274,6 +274,10 @@ def diff(a_tensor, n=1, axis=-1, prepend_tensor=None, append_tensor=None):
     if n < 0:
         raise ValueError(f"order must be non-negative but got {n}")
 
+    if n == 0:
+        # match numpy and return the input immediately
+        return a_tensor
+
     if prepend_tensor is not None:
         shape = list(a_tensor.shape)
         shape[axis] = prepend_tensor.shape[axis] if prepend_tensor.ndim > 0 else 1
@@ -357,6 +361,14 @@ 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
 
 
@@ -450,6 +462,10 @@ def indices(dimensions, dtype=int, sparse=False):
 
 
 def bincount(x_tensor, /, weights_tensor=None, minlength=0):
+    if x_tensor.numel() == 0:
+        # edge case allowed by numpy
+        x_tensor = torch.as_tensor([], dtype=int)
+
     int_dtype = _dtypes_impl.default_int_dtype
     (x_tensor,) = _util.cast_dont_broadcast((x_tensor,), int_dtype, casting="safe")
 
@@ -460,6 +476,14 @@ def bincount(x_tensor, /, weights_tensor=None, minlength=0):
 # ### linspace, geomspace, logspace and arange ###
 
 
+def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
+    if axis != 0 or retstep or not endpoint:
+        raise NotImplementedError
+    # XXX: raises TypeError if start or stop are not scalars
+    result = torch.linspace(start, stop, num, dtype=dtype)
+    return result
+
+
 def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     if axis != 0 or not endpoint:
         raise NotImplementedError
@@ -474,6 +498,13 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     return result
 
 
+def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
+    if axis != 0 or not endpoint:
+        raise NotImplementedError
+    result = torch.logspace(start, stop, num, base=base, dtype=dtype)
+    return result
+
+
 def arange(start=None, stop=None, step=1, dtype=None):
     if step == 0:
         raise ZeroDivisionError
@@ -523,36 +554,75 @@ def eye(N, M=None, k=0, dtype=float):
     return z
 
 
-def zeros_like(a, dtype=None, shape=None):
+def zeros(shape, dtype=None, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    if dtype is None:
+        dtype = _dtypes_impl.default_float_dtype
+    result = torch.zeros(shape, dtype=dtype)
+    return result
+
+
+def zeros_like(a, dtype=None, shape=None, order="K"):
+    if order != "K":
+        raise NotImplementedError
     result = torch.zeros_like(a, dtype=dtype)
     if shape is not None:
         result = result.reshape(shape)
     return result
 
 
-def ones_like(a, dtype=None, shape=None):
+def ones(shape, dtype=None, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    if dtype is None:
+        dtype = _dtypes_impl.default_float_dtype
+    result = torch.ones(shape, dtype=dtype)
+    return result
+
+
+def ones_like(a, dtype=None, shape=None, order="K"):
+    if order != "K":
+        raise NotImplementedError
     result = torch.ones_like(a, dtype=dtype)
     if shape is not None:
         result = result.reshape(shape)
     return result
 
 
-def full_like(a, fill_value, dtype=None, shape=None):
+def full_like(a, fill_value, dtype=None, shape=None, order="K"):
+    if order != "K":
+        raise NotImplementedError
     # XXX: fill_value broadcasts
     result = torch.full_like(a, fill_value, dtype=dtype)
     if shape is not None:
         result = result.reshape(shape)
     return result
 
 
-def empty_like(prototype, dtype=None, shape=None):
+def empty(shape, dtype=None, order="C"):
+    if order != "C":
+        raise NotImplementedError
+    if dtype is None:
+        dtype = _dtypes_impl.default_float_dtype
+    result = torch.empty(shape, dtype=dtype)
+    return result
+
+
+def empty_like(prototype, dtype=None, shape=None, order="K"):
+    if order != "K":
+        raise NotImplementedError
     result = torch.empty_like(prototype, dtype=dtype)
     if shape is not None:
         result = result.reshape(shape)
     return result
 
 
-def full(shape, fill_value, dtype=None):
+def full(shape, fill_value, dtype=None, order="C"):
+    if isinstance(shape, int):
+        shape = (shape,)
+    if order != "C":
+        raise NotImplementedError
     if dtype is None:
         dtype = fill_value.dtype
     if not isinstance(shape, (tuple, list)):

torch_np/_funcs.py

@@ -50,6 +50,12 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis=None):
     return _helpers.array_from(result)
 
 
+@normalizer
+def tile(A: ArrayLike, reps):
+    result = _impl.tile(A, reps)
+    return _helpers.array_from(result)
+
+
 # ### diag et al ###
 
 
@@ -449,8 +455,45 @@ def quantile(
     *,
     interpolation=None,
 ):
-    if interpolation is not None:
-        raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
+    result = _impl.quantile(
+        a,
+        q,
+        axis,
+        overwrite_input=overwrite_input,
+        method=method,
+        keepdims=keepdims,
+        interpolation=interpolation,
+    )
+    return _helpers.result_or_out(result, out, promote_scalar=True)
 
-    result = _impl.quantile(a, q, axis, method=method, keepdims=keepdims)
+
+@normalizer
+def percentile(
+    a: ArrayLike,
+    q: ArrayLike,
+    axis: AxisLike = None,
+    out: Optional[NDArray] = None,
+    overwrite_input=False,
+    method="linear",
+    keepdims=False,
+    *,
+    interpolation=None,
+):
+    result = _impl.percentile(
+        a,
+        q,
+        axis,
+        overwrite_input=overwrite_input,
+        method=method,
+        keepdims=keepdims,
+        interpolation=interpolation,
+    )
     return _helpers.result_or_out(result, out, promote_scalar=True)
+
+
+def median(
+    a, axis=None, out: Optional[NDArray] = None, overwrite_input=False, keepdims=False
+):
+    return quantile(
+        a, 0.5, axis=axis, overwrite_input=overwrite_input, out=out, keepdims=keepdims
+    )