linting & cleanup

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -938,13 +938,16 @@ def quantile(
         9. 'normal_unbiased'
 
         The first three methods are discontinuous.
-        Only 'linear' is implemented for now.
+        Only 'linear', 'inverted_cdf' and 'averaged_inverted_cdf' are implemented.
 
     keepdims : bool, optional
         If this is set to True, the axes which are reduced are left in
         the result as dimensions with size one. With this option, the
         result will broadcast correctly against the original array `a`.
 
+    nan_policy : str, optional
+        'propagate' (default) or 'omit'.
+
     weights : array_like, optional
         An array of weights associated with the values in `a`. Each value in
         `a` contributes to the quantile according to its associated weight.
@@ -1121,20 +1124,24 @@ def quantile(
         msg = "`q` values must be in the range [0, 1]"
         raise ValueError(msg)
 
-    # Delegate where possible.
+    # Delegate when possible.
     if is_numpy_namespace(xp) and nan_policy == "propagate":
+        # TODO: call nanquantile for nan_policy == "omit" once
+        # https://github.com/numpy/numpy/issues/29709 is fixed
         return xp.quantile(
             a, q_arr, axis=axis, method=method, keepdims=keepdims, weights=weights
         )
-    # No delegation for dask: I couldn't make it work
-    basic_case = method == "linear" and weights is None and nan_policy == "propagate"
-    if (basic_case and is_jax_namespace(xp)) or is_cupy_namespace(xp):
+    # No delegation for dask: I couldn't make it work.
+    basic_case = method == "linear" and weights is None
+    jax_or_cupy = is_jax_namespace(xp) or is_cupy_namespace(xp)
+    if basic_case and nan_policy == "propagate" and jax_or_cupy:
         return xp.quantile(a, q_arr, axis=axis, method=method, keepdims=keepdims)
     if basic_case and is_torch_namespace(xp):
-        return xp.quantile(a, q_arr, dim=axis, interpolation=method, keepdim=keepdims)
+        quantile = xp.quantile if nan_policy == "propagate" else xp.nanquantile
+        return quantile(a, q_arr, dim=axis, interpolation=method, keepdim=keepdims)
 
-    # XXX: I'm not sure we want to support dask, it seems uterly slow...
     # Otherwise call our implementation (will sort data)
+    # XXX: I'm not sure we want to support dask, it seems uterly slow...
     return _quantile.quantile(
         a,
         q_arr,

src/array_api_extra/_lib/_quantile.py

@@ -73,9 +73,10 @@ def quantile(  # numpydoc ignore=PR01,RT01
     return res[0, ...] if q_scalar else res
 
 
-def _quantile(  # numpydoc ignore=GL08
+def _quantile(  # numpydoc ignore=PR01,RT01
     a: Array, q: Array, n: int, axis: int, method: str, xp: ModuleType
 ) -> Array:
+    """Compute quantile by sorting `a`."""
     a = xp.sort(a, axis=axis, stable=False)
     mask_nan = xp.any(xp.isnan(a), axis=axis, keepdims=True)
     if xp.any(mask_nan):
@@ -114,7 +115,7 @@ def _quantile(  # numpydoc ignore=GL08
     )
 
 
-def _weighted_quantile(
+def _weighted_quantile(  # numpydoc ignore=PR01,RT01
     a: Array,
     q: Array,
     weights: Array,
@@ -126,7 +127,9 @@ def _weighted_quantile(
     device: Device,
 ) -> Array:
     """
-    a is expected to be 1d or 2d.
+    Compute weighted quantile using searchsorted on CDF.
+
+    `a` is expected to be 1d or 2d.
     """
     a = xp.moveaxis(a, axis, -1)
     if weights.ndim > 1:
@@ -151,7 +154,7 @@ def _weighted_quantile(
     return xp.stack(res, axis=1)
 
 
-def _weighted_quantile_sorted_1d(
+def _weighted_quantile_sorted_1d(  # numpydoc ignore=GL08
     x: Array,
     q: Array,
     w: Array,
@@ -165,10 +168,10 @@ def _weighted_quantile_sorted_1d(
         w = xp.where(xp.isnan(x), 0.0, w)
     elif xp.any(xp.isnan(x)):
         return xp.full(q.shape, xp.nan, dtype=x.dtype, device=device)
-    cw = xp.cumulative_sum(w)
-    t = cw[-1] * q
-    i = xp.searchsorted(cw, t, side="left")
-    j = xp.searchsorted(cw, t, side="right")
+    cdf = xp.cumulative_sum(w)
+    t = cdf[-1] * q
+    i = xp.searchsorted(cdf, t, side="left")
+    j = xp.searchsorted(cdf, t, side="right")
     i = xp.clip(i, 0, n - 1)
     j = xp.clip(j, 0, n - 1)