untested implem; limited to method="linear"; trying to mimic numpy behavior

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -910,16 +910,29 @@ def quantile(
     """
     TODO
     """
-
     methods = {"linear"}
+
     if method not in methods:
         message = f"`method` must be one of {methods}"
         raise ValueError(message)
+    if keepdims not in {True, False}:
+        message = "If specified, `keepdims` must be True or False."
+        raise ValueError(message)
     if xp is None:
         xp = array_namespace(a)
-    if a.ndim < 1:
+
+    a = xp.asarray(a)
+    if not xp.isdtype(a.dtype, ('integral', 'real floating')):
+        raise ValueError("`a` must have real dtype.")
+    if not xp.isdtype(xp.asarray(q).dtype, 'real floating'):
+        raise ValueError("`q` must have real floating dtype.")
+    ndim = a.ndim
+    if ndim < 1:
         msg = "`a` must be at least 1-dimensional"
         raise TypeError(msg)
+    if (axis >= ndim) or (axis < -ndim):
+        message = "`axis` is not compatible with the dimension of `a`."
+        raise ValueError(message)
 
     # Delegate where possible.
     if is_numpy_namespace(xp) or is_dask_namespace(xp):

src/array_api_extra/_lib/_quantile.py

@@ -10,8 +10,8 @@
 
 
 def quantile(  # numpydoc ignore=PR01,RT01
-    x,
-    p,
+    a: Array,
+    q: Array | float,
     /,
     method: str = 'linear',  # noqa: ARG001
     axis: int | None = None,
@@ -20,107 +20,67 @@ def quantile(  # numpydoc ignore=PR01,RT01
     xp: ModuleType,
 ):
     """See docstring in `array_api_extra._delegation.py`."""
-    # Input validation / standardization
-    temp = _quantile_iv(x, p, axis, keepdims)
-    y, p, axis, keepdims, n, axis_none, ndim = temp
+    device = get_device(a)
+    floating_dtype = xp.result_type(a, xp.asarray(q))
+    a = xp.asarray(a, dtype=floating_dtype, device=device)
+    q = xp.asarray(q, dtype=floating_dtype, device=device)
 
-    res = _quantile_hf(y, p, n, xp)
+    if xp.any((q > 1) | (q < 0) | xp.isnan(q)):
+        raise ValueError("`q` values must be in the range [0, 1]")
 
-    # Reshape per axis/keepdims
-    if axis_none and keepdims:
-        shape = (1,)*(ndim - 1) + res.shape
-        res = xp.reshape(res, shape)
-        axis = -1
-
-    res = xp.moveaxis(res, -1, axis)
-
-    if not keepdims:
-        res = xp.squeeze(res, axis=axis)
-
-    return res[()] if res.ndim == 0 else res
-
-
-def _quantile_iv(
-    x: Array,
-    p: Array,
-    axis: int | None,
-    keepdims: bool,
-    xp: ModuleType
-):
-
-    if not xp.isdtype(xp.asarray(x).dtype, ('integral', 'real floating')):
-        raise ValueError("`x` must have real dtype.")
-
-    if not xp.isdtype(xp.asarray(p).dtype, 'real floating'):
-        raise ValueError("`p` must have real floating dtype.")
-
-    p_mask = (p > 1) | (p < 0) | xp.isnan(p)
-    if xp.any(p_mask):
-        raise ValueError("`p` values must be in the range [0, 1]")
-
-    device = get_device(x)
-    floating_dtype = xp.result_type(x, p)
-    x = xp.asarray(x, dtype=floating_dtype, device=device)
-    p = xp.asarray(p, dtype=floating_dtype, device=device)
-    dtype = x.dtype
+    q_scalar = q.ndim == 0
+    if q_scalar:
+        q = xp.reshape(q, (1,))
 
     axis_none = axis is None
-    ndim = max(x.ndim, p.ndim)
     if axis_none:
-        x = xp.reshape(x, (-1,))
-        p = xp.reshape(p, (-1,))
+        a = xp.reshape(a, (-1,))
         axis = 0
-    elif np.iterable(axis) or int(axis) != axis:
-        message = "`axis` must be an integer or None."
-        raise ValueError(message)
-    elif (axis >= ndim) or (axis < -ndim):
-        message = "`axis` is not compatible with the shapes of the inputs."
-        raise ValueError(message)
     axis = int(axis)
 
-    if keepdims not in {None, True, False}:
-        message = "If specified, `keepdims` must be True or False."
-        raise ValueError(message)
-
+    n = eager_shape(a, axis)
     # If data has length zero along `axis`, the result will be an array of NaNs just
     # as if the data had length 1 along axis and were filled with NaNs.
-    n = eager_shape(x, axis)
     if n == 0:
-        shape = eager_shape(x)
+        shape = list(eager_shape(a))
         shape[axis] = 1
         n = 1
-        x = xp.full(shape, xp.nan, dtype=dtype, device=device)
-
-    y = xp.sort(x, axis=axis, stable=False)
-    # FIXME: I still need to look into the broadcasting:
-    y, p = _broadcast_arrays((y, p), axis=axis)
+        a = xp.full(shape, xp.nan, dtype=floating_dtype, device=device)
 
-    p_shape = eager_shape(p)
-    if (keepdims is False) and (p_shape[axis] != 1):
-        message = "`keepdims` may be False only if the length of `p` along `axis` is 1."
-        raise ValueError(message)
-    keepdims = (p_shape[axis] != 1) if keepdims is None else keepdims
+    a = xp.sort(a, axis=axis, stable=False)
+    # to support weights, the main thing would be to
+    # argsort a, and then use it to sort a and w.
+    # The hard part will be dealing with 0-weights and NaNs
+    # But maybe a proper use of searchsorted + left/right side will work?
 
-    y = xp.moveaxis(y, axis, -1)
-    p = xp.moveaxis(p, axis, -1)
+    res = _quantile_hf(a, q, n, axis, xp)
 
-    nans = xp.isnan(y)
-    nan_out = xp.any(nans, axis=-1)
-    if xp.any(nan_out):
-        y = xp.asarray(y, copy=True)  # ensure writable
-        y = at(y, nan_out).set(xp.nan)
+    # reshaping to conform to doc/other libs' behavior
+    if axis_none:
+        if keepdims:
+            res = xp.reshape(res, q.shape + (1,) * a.ndim)
+    else:
+        res = xp.moveaxis(res, axis, 0)
+        if keepdims:
+            shape = list(a.shape)
+            shape[axis] = 1
+            shape = q.shape + tuple(shape)
+            res = xp.reshape(res, shape)
 
-    return y, p, axis, keepdims, n, axis_none, ndim, xp
+    return res[0, ...] if q_scalar else res
 
 
-def _quantile_hf(y, p, n, xp):
+def _quantile_hf(y: Array, p: Array, n: int, axis: int, xp: ModuleType):
     m = 1 - p
     jg = p*n + m - 1
     j = jg // 1
     g = jg % 1
     g[j < 0] = 0
     j = xp.clip(j, 0., n - 1)
     jp1 = xp.clip(j + 1, 0., n - 1)
+    # `̀j` and `jp1` are 1d arrays
 
-    return ((1 - g) * xp.take_along_axis(y, xp.astype(j, xp.int64), axis=-1)
-            + g * xp.take_along_axis(y, xp.astype(jp1, xp.int64), axis=-1))
+    return (
+        (1 - g) * xp.take(y, xp.astype(j, xp.int64), axis=axis)
+        + g * xp.take(y, xp.astype(jp1, xp.int64), axis=axis)
+    )