Merge pull request #29 from Quansight-Labs/function_base

finish up reductions and related animals

Author: ev-br

autogen/numpy_api_dump.py

@@ -242,14 +242,6 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
     raise NotImplementedError
 
 
-def cumprod(a, axis=None, dtype=None, out=None):
-    raise NotImplementedError
-
-
-def cumproduct(*args, **kwargs):
-    raise NotImplementedError
-
-
 def cumsum(a, axis=None, dtype=None, out=None):
     raise NotImplementedError
 
@@ -610,10 +602,6 @@ def may_share_memory(a, b, /, max_work=None):
     raise NotImplementedError
 
 
-def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
-    raise NotImplementedError
-
-
 def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
     raise NotImplementedError
 
@@ -734,20 +722,6 @@ def partition(a, kth, axis=-1, kind="introselect", order=None):
     raise NotImplementedError
 
 
-def percentile(
-    a,
-    q,
-    axis=None,
-    out=None,
-    overwrite_input=False,
-    method="linear",
-    keepdims=False,
-    *,
-    interpolation=None,
-):
-    raise NotImplementedError
-
-
 def piecewise(x, condlist, funclist, *args, **kw):
     raise NotImplementedError
 
@@ -800,10 +774,6 @@ def product(*args, **kwargs):
     raise NotImplementedError
 
 
-def ptp(a, axis=None, out=None, keepdims=NoValue):
-    raise NotImplementedError
-
-
 def put(a, ind, v, mode="raise"):
     raise NotImplementedError
 
@@ -816,20 +786,6 @@ def putmask(a, mask, values):
     raise NotImplementedError
 
 
-def quantile(
-    a,
-    q,
-    axis=None,
-    out=None,
-    overwrite_input=False,
-    method="linear",
-    keepdims=False,
-    *,
-    interpolation=None,
-):
-    raise NotImplementedError
-
-
 def ravel(a, order="C"):
     raise NotImplementedError
 

torch_np/__init__.py

@@ -10,6 +10,8 @@
 
 # from . import testing
 
+alltrue = all
+sometrue = any
 
 inf = float("inf")
 nan = float("nan")

torch_np/_decorators.py

@@ -93,10 +93,7 @@ def axis_keepdims_wrapper(func):
     see them). The `axis` argument we normalize and pass through to pytorch functions.
 
     """
-    # XXX: move this out of _ndarray.py (circular imports)
-    #
-    # TODO: 1. get rid of _helpers.result_or_out
-    #       2. sort out function signatures: how they flow through all decorators etc
+    # TODO: sort out function signatures: how they flow through all decorators etc
     @functools.wraps(func)
     def wrapped(a, axis=None, keepdims=NoValue, *args, **kwds):
         from ._ndarray import asarray, ndarray
@@ -107,7 +104,36 @@ def wrapped(a, axis=None, keepdims=NoValue, *args, **kwds):
         if isinstance(axis, ndarray):
             axis = operator.index(axis)
 
-        result = _util.axis_keepdims(func, tensor, axis, keepdims, *args, **kwds)
+        result = _util.axis_expand_func(func, tensor, axis, *args, **kwds)
+
+        if keepdims:
+            result = _util.apply_keepdims(result, axis, tensor.ndim)
+
+        return result
+
+    return wrapped
+
+
+def axis_none_ravel_wrapper(func):
+    """`func` accepts an array-like as a 1st arg, returns a tensor.
+
+    This decorator implements the generic handling of axis=None acting on a
+    raveled array. One use is cumprod / cumsum. concatenate also uses a
+    similar logic.
+
+    """
+
+    @functools.wraps(func)
+    def wrapped(a, axis=None, *args, **kwds):
+        from ._ndarray import asarray, ndarray
+
+        tensor = asarray(a).get()
+
+        # standardize the axis argument
+        if isinstance(axis, ndarray):
+            axis = operator.index(axis)
+
+        result = _util.axis_ravel_func(func, tensor, axis, *args, **kwds)
         return result
 
     return wrapped

torch_np/_detail/_reductions.py

@@ -89,6 +89,11 @@ def min(tensor, axis=None, initial=NoValue, where=NoValue):
     return result
 
 
+def ptp(tensor, axis=None):
+    result = tensor.amax(axis) - tensor.amin(axis)
+    return result
+
+
 def sum(tensor, axis=None, dtype=None, initial=NoValue, where=NoValue):
     if initial is not NoValue or where is not NoValue:
         raise NotImplementedError
@@ -161,3 +166,97 @@ def var(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
     result = tensor.var(dim=axis, correction=ddof)
 
     return result
+
+
+# cumsum / cumprod are almost reductions:
+#   1. no keepdims
+#   2. axis=None ravels (cf concatenate)
+
+
+def cumprod(tensor, axis, dtype=None):
+    if dtype == torch.bool:
+        dtype = _scalar_types.default_int_type.dtype
+    if dtype is None:
+        dtype = tensor.dtype
+
+    result = tensor.cumprod(axis=axis, dtype=dtype)
+
+    return result
+
+
+def cumsum(tensor, axis, dtype=None):
+    if dtype == torch.bool:
+        dtype = _scalar_types.default_int_type.dtype
+    if dtype is None:
+        dtype = tensor.dtype
+
+    result = tensor.cumsum(axis=axis, dtype=dtype)
+
+    return result
+
+
+def average(a_tensor, axis, w_tensor):
+
+    # dtype
+    # FIXME: 1. use result_type
+    #        2. actually implement multiply w/dtype
+    if not a_tensor.dtype.is_floating_point:
+        result_dtype = torch.float64
+        a_tensor = a_tensor.to(result_dtype)
+
+    # axis
+    if axis is None:
+        (a_tensor, w_tensor), axis = _util.axis_none_ravel(
+            a_tensor, w_tensor, axis=axis
+        )
+
+    # axis & weights
+    if a_tensor.shape != w_tensor.shape:
+        if axis is None:
+            raise TypeError(
+                "Axis must be specified when shapes of a and weights " "differ."
+            )
+        if w_tensor.ndim != 1:
+            raise TypeError("1D weights expected when shapes of a and weights differ.")
+        if w_tensor.shape[0] != a_tensor.shape[axis]:
+            raise ValueError("Length of weights not compatible with specified axis.")
+
+        # setup weight to broadcast along axis
+        w_tensor = torch.broadcast_to(
+            w_tensor, (a_tensor.ndim - 1) * (1,) + w_tensor.shape
+        )
+        w_tensor = w_tensor.swapaxes(-1, axis)
+
+    # do the work
+    numerator = torch.mul(a_tensor, w_tensor).sum(axis)
+    denominator = w_tensor.sum(axis)
+    result = numerator / denominator
+
+    return result, denominator
+
+
+def quantile(a_tensor, q_tensor, axis, method):
+
+    if (0 > q_tensor).any() or (q_tensor > 1).any():
+        raise ValueError("Quantiles must be in range [0, 1], got %s" % q_tensor)
+
+    if not a_tensor.dtype.is_floating_point:
+        dtype = _scalar_types.default_float_type.torch_dtype
+        a_tensor = a_tensor.to(dtype)
+
+    # edge case: torch.quantile only supports float32 and float64
+    if a_tensor.dtype == torch.float16:
+        a_tensor = a_tensor.to(torch.float32)
+
+    # axis
+    if axis is not None:
+        axis = _util.normalize_axis_tuple(axis, a_tensor.ndim)
+    axis = _util.allow_only_single_axis(axis)
+
+    q_tensor = q_tensor.to(a_tensor.dtype)
+
+    (a_tensor, q_tensor), axis = _util.axis_none_ravel(a_tensor, q_tensor, axis=axis)
+
+    result = torch.quantile(a_tensor, q_tensor, axis=axis, interpolation=method)
+
+    return result

torch_np/_detail/_util.py

@@ -111,7 +111,8 @@ def expand_shape(arr_shape, axis):
 def apply_keepdims(tensor, axis, ndim):
     if axis is None:
         # tensor was a scalar
-        tensor = torch.full((1,) * ndim, fill_value=tensor, dtype=tensor.dtype)
+        shape = (1,) * ndim
+        tensor = tensor.expand(shape).contiguous()  # avoid CUDA synchronization
     else:
         shape = expand_shape(tensor.shape, axis)
         tensor = tensor.reshape(shape)
@@ -211,8 +212,8 @@ def cast_and_broadcast(tensors, out_param, casting):
     return tuple(processed_tensors)
 
 
-def axis_keepdims(func, tensor, axis, keepdims, *args, **kwds):
-    """Generically handle axis and keepdims arguments in reductions."""
+def axis_expand_func(func, tensor, axis, *args, **kwds):
+    """Generically handle axis arguments in reductions."""
     if axis is not None:
         if not isinstance(axis, (list, tuple)):
             axis = (axis,)
@@ -225,8 +226,18 @@ def axis_keepdims(func, tensor, axis, keepdims, *args, **kwds):
 
     result = func(tensor, axis=axis, *args, **kwds)
 
-    if keepdims:
-        result = apply_keepdims(result, axis, tensor.ndim)
+    return result
+
+
+def axis_ravel_func(func, tensor, axis, *args, **kwds):
+    """Generically handle axis arguments in cumsum/cumprod."""
+    if axis is not None:
+        axis = normalize_axis_index(axis, tensor.ndim)
+
+    tensors, axis = axis_none_ravel(tensor, axis=axis)
+    tensor = tensors[0]
+
+    result = func(tensor, axis=axis, *args, **kwds)
 
     return result
 

torch_np/_helpers.py

@@ -40,13 +40,26 @@ def cast_and_broadcast(tensors, out, casting):
     return tuple(tensors)
 
 
-def result_or_out(result_tensor, out_array=None):
-    """A helper for returns with out= argument."""
+def result_or_out(result_tensor, out_array=None, promote_scalar=False):
+    """A helper for returns with out= argument.
+
+    If `promote_scalar is True`, then:
+        if result_tensor.numel() == 1 and out is zero-dimensional,
+            result_tensor is placed into the out array.
+    This weirdness is used e.g. in `np.percentile`
+    """
     if out_array is not None:
         if not isinstance(out_array, ndarray):
             raise TypeError("Return arrays must be of ArrayType")
         if result_tensor.shape != out_array.shape:
-            raise ValueError("Bad size of the out array.")
+            can_fit = result_tensor.numel() == 1 and out_array.ndim == 0
+            if promote_scalar and can_fit:
+                result_tensor = result_tensor.squeeze()
+            else:
+                raise ValueError(
+                    f"Bad size of the out array: out.shape = {out_array.shape}"
+                    f" while result.shape = {result_tensor.shape}."
+                )
         out_tensor = out_array.get()
         out_tensor.copy_(result_tensor)
         return out_array

torch_np/_ndarray.py

@@ -3,7 +3,13 @@
 import torch
 
 from . import _binary_ufuncs, _dtypes, _helpers, _unary_ufuncs
-from ._decorators import NoValue, axis_keepdims_wrapper, dtype_to_torch, emulate_out_arg
+from ._decorators import (
+    NoValue,
+    axis_keepdims_wrapper,
+    axis_none_ravel_wrapper,
+    dtype_to_torch,
+    emulate_out_arg,
+)
 from ._detail import _reductions, _util
 
 newaxis = None
@@ -274,13 +280,21 @@ def nonzero(self):
     all = emulate_out_arg(axis_keepdims_wrapper(_reductions.all))
     max = emulate_out_arg(axis_keepdims_wrapper(_reductions.max))
     min = emulate_out_arg(axis_keepdims_wrapper(_reductions.min))
+    ptp = emulate_out_arg(axis_keepdims_wrapper(_reductions.ptp))
 
     sum = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.sum)))
     prod = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.prod)))
     mean = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.mean)))
     var = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.var)))
     std = emulate_out_arg(axis_keepdims_wrapper(dtype_to_torch(_reductions.std)))
 
+    cumprod = emulate_out_arg(
+        axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumprod))
+    )
+    cumsum = emulate_out_arg(
+        axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumsum))
+    )
+
     ### indexing ###
     def __getitem__(self, *args, **kwds):
         t_args = _helpers.ndarrays_to_tensors(*args)
@@ -380,6 +394,8 @@ def result_type(*arrays_and_dtypes):
             dtypes.append(_dtypes.dtype(entry))
         elif isinstance(entry, _dtypes.DType):
             dtypes.append(entry)
+        elif isinstance(entry, str):
+            dtypes.append(_dtypes.dtype(entry))
         else:
             dtypes.append(asarray(entry).dtype)