MAINT: quantile/percentile/median

Deviate from the standard argument handling:
- percentile and median delegate to quantile at the wrapper level; this is just less code
- keepdims=True handling is inline in _reductions.py::quantile, not in a decorator.
  The standard decorator expects the axis as the second argument, while here it is the
  third one. Can be fixed, but seems to be more hassle then worth TBH.

Author: ev-br

torch_np/_detail/_reductions.py

@@ -337,7 +337,7 @@ def average(a_tensor, axis, w_tensor):
     return result, denominator
 
 
-def quantile(a_tensor, q_tensor, axis, method):
+def quantile(a_tensor, q_tensor, axis, method, keepdims=False):
 
     if (0 > q_tensor).any() or (q_tensor > 1).any():
         raise ValueError("Quantiles must be in range [0, 1], got %s" % q_tensor)
@@ -350,15 +350,24 @@ def quantile(a_tensor, q_tensor, axis, method):
     if a_tensor.dtype == torch.float16:
         a_tensor = a_tensor.to(torch.float32)
 
+    # TODO: consider moving this normalize_axis_tuple dance to normalize axis? Across the board if at all.
     # axis
     if axis is not None:
         axis = _util.normalize_axis_tuple(axis, a_tensor.ndim)
     axis = _util.allow_only_single_axis(axis)
 
     q_tensor = _util.cast_if_needed(q_tensor, a_tensor.dtype)
 
+
+    # axis=None ravels, so store the originals to reuse with keepdims=True below
+    ax, ndim = axis, a_tensor.ndim
     (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)
 
+    # NB: not using @emulate_keepdims here because the signature is (a, q, axis, ...)
+    # while the decorator expects (a, axis, ...)
+    # this can be fixed, of course, but the cure seems worse then the desease
+    if keepdims:
+        result = _util.apply_keepdims(result, ax, ndim)
     return result

torch_np/_funcs.py

@@ -459,3 +459,22 @@ def cumsum(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=Non
 def cumprod(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=None):
     result = _reductions.cumprod(a, axis=axis, dtype=dtype)
     return _helpers.result_or_out(result, out)
+
+
+@normalizer
+def quantile(
+    a : ArrayLike,
+    q : ArrayLike,
+    axis: AxisLike=None,
+    out=None,
+    overwrite_input=False,
+    method="linear",
+    keepdims=False,
+    *,
+    interpolation=None,
+):
+    if interpolation is not None:
+        raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
+
+    result = _reductions.quantile(a, q, axis, method=method, keepdims=keepdims)
+    return _helpers.result_or_out(result, out, promote_scalar=True)

torch_np/_wrapper.py

@@ -674,6 +674,7 @@ def average(a, axis=None, weights=None, returned=False, *, keepdims=NoValue):
         return _helpers.array_from(result)
 
 
+# Normalizations (ArrayLike et al) are done in quantile.
 def percentile(
     a,
     q,
@@ -685,36 +686,14 @@ def percentile(
     *,
     interpolation=None,
 ):
-    return quantile(
+    return _funcs.quantile(
         a, asarray(q) / 100.0, axis, out, overwrite_input, method, keepdims=keepdims
     )
 
 
-def quantile(
-    a,
-    q,
-    axis=None,
-    out=None,
-    overwrite_input=False,
-    method="linear",
-    keepdims=False,
-    *,
-    interpolation=None,
-):
-    if interpolation is not None:
-        raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
-
-    a_tensor, q_tensor = _helpers.to_tensors(a, q)
-    result = _reductions.quantile(a_tensor, q_tensor, axis, method)
-
-    # keepdims
-    if keepdims:
-        result = _util.apply_keepdims(result, axis, a_tensor.ndim)
-    return _helpers.result_or_out(result, out, promote_scalar=True)
-
 
 def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
-    return quantile(
+    return _funcs.quantile(
         a, 0.5, axis=axis, overwrite_input=overwrite_input, out=out, keepdims=keepdims
     )