Add weighted quantile and percentile support with tests

Author: Aniketsy

jax/_src/numpy/reductions.py

@@ -2337,7 +2337,8 @@ def cumulative_prod(
 @partial(api.jit, static_argnames=('axis', 'overwrite_input', 'interpolation', 'keepdims', 'method'))
 def quantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = None,
              out: None = None, overwrite_input: bool = False, method: str = "linear",
-             keepdims: bool = False, *, interpolation: DeprecatedArg | str = DeprecatedArg()) -> Array:
+             keepdims: bool = False, weights: ArrayLike | None = None, *,
+             interpolation: DeprecatedArg | str = DeprecatedArg()) -> Array:
   """Compute the quantile of the data along the specified axis.
 
   JAX implementation of :func:`numpy.quantile`.
@@ -2387,7 +2388,7 @@ def quantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = No
       ("The interpolation= argument to 'quantile' is deprecated. "
        "Use 'method=' instead."), stacklevel=2)
     method = interpolation
-  return _quantile(lax_internal.asarray(a), lax_internal.asarray(q), axis, method, keepdims, False)
+  return _quantile(lax_internal.asarray(a), lax_internal.asarray(q), axis, method, keepdims, False, weights)
 
 # TODO(jakevdp): interpolation argument deprecated 2024-05-16
 @export
@@ -2449,7 +2450,7 @@ def nanquantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None =
   return _quantile(lax_internal.asarray(a), lax_internal.asarray(q), axis, method, keepdims, True)
 
 def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
-              method: str, keepdims: bool, squash_nans: bool) -> Array:
+              method: str, keepdims: bool, squash_nans: bool, weights: ArrayLike | None = None) -> Array:
   if method not in ["linear", "lower", "higher", "midpoint", "nearest"]:
     raise ValueError("method can only be 'linear', 'lower', 'higher', 'midpoint', or 'nearest'")
   a, = promote_dtypes_inexact(a)
@@ -2488,6 +2489,66 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     raise ValueError(f"q must be have rank <= 1, got shape {q.shape}")
 
   a_shape = a.shape
+  # Handle weights
+  if weights is not None:
+    a, weights = promote_dtypes_inexact(a, weights)
+    if axis is None:
+      a = a.ravel()
+      weights = weights.ravel()
+      axis = 0
+    else:
+      weights = _broadcast_to(weights, a.shape)
+    if squash_nans:
+      nan_mask = ~lax_internal._isnan(a)
+      if axis is None:
+        a = a[nan_mask]
+        weights = weights[nan_mask]
+      else:
+        weights = _where(nan_mask, weights, 0)
+    a_sorted, weights_sorted = lax.sort_key_val(a, weights, dimension=axis)
+
+    cum_weights = lax.cumsum(weights_sorted, axis=axis)
+    total_weight = lax.sum(weights_sorted, axis=axis, keepdims=True)
+    if lax_internal._all(total_weight == 0):
+      raise ValueError("Sum of weights must not be zero.")
+    cum_weights_norm = cum_weights / total_weight
+    quantile_pos = q
+    mask = cum_weights_norm >= quantile_pos[..., None]
+    idx = lax.argmin(mask.astype(int), axis=axis)
+    idx_prev = lax.max(idx - 1, _lax_const(idx, 0))
+    idx_next = idx
+    gather_shape = list(a_sorted.shape)
+    gather_shape[axis] = 1
+    dnums = lax.GatherDimensionNumbers(
+      offset_dims=tuple(range(len(a_sorted.shape))),
+      collapsed_slice_dims=(axis,),
+      start_index_map=(axis,))
+    prev_value = lax.gather(a_sorted, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=gather_shape)
+    next_value = lax.gather(a_sorted, idx_next[..., None], dimension_numbers=dnums, slice_sizes=gather_shape)
+    prev_cumw = lax.gather(cum_weights_norm, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=gather_shape)
+    next_cumw = lax.gather(cum_weights_norm, idx_next[..., None], dimension_numbers=dnums, slice_sizes=gather_shape)
+
+    if method == "linear":
+      denom = next_cumw - prev_cumw
+      denom = lax.select(denom == 0, _lax_const(denom, 1), denom)
+      weight = (quantile_pos - prev_cumw) / denom
+      result = prev_value * (1 - weight) + next_value * weight
+    elif method == "lower":
+      result = prev_value
+    elif method == "higher":
+      result = next_value
+    elif method == "nearest":
+      use_prev = (quantile_pos - prev_cumw) < (next_cumw - quantile_pos)
+      result = lax.select(use_prev, prev_value, next_value)
+    elif method == "midpoint":
+      result = (prev_value + next_value) / 2
+    else:
+      raise ValueError(f"{method=!r} not recognized")
+
+    if not keepdims:
+      result = lax.squeeze(result, axis)
+    return lax.convert_element_type(result, a.dtype)
+
 
   if squash_nans:
     a = _where(lax_internal._isnan(a), np.nan, a) # Ensure nans are positive so they sort to the end.
@@ -2578,7 +2639,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
 def percentile(a: ArrayLike, q: ArrayLike,
                axis: int | tuple[int, ...] | None = None,
                out: None = None, overwrite_input: bool = False, method: str = "linear",
-               keepdims: bool = False, *, interpolation: str | DeprecatedArg = DeprecatedArg()) -> Array:
+               keepdims: bool = False, weights: ArrayLike | None = None, *, interpolation: str | DeprecatedArg = DeprecatedArg()) -> Array:
   """Compute the percentile of the data along the specified axis.
 
   JAX implementation of :func:`numpy.percentile`.
@@ -2627,7 +2688,7 @@ def percentile(a: ArrayLike, q: ArrayLike,
        "Use 'method=' instead."), stacklevel=2)
     method = interpolation
   return quantile(a, q / 100, axis=axis, out=out, overwrite_input=overwrite_input,
-                  method=method, keepdims=keepdims)
+                  method=method, keepdims=keepdims, weights=weights)
 
 
 # TODO(jakevdp): interpolation argument deprecated 2024-05-16
@@ -2636,7 +2697,7 @@ def percentile(a: ArrayLike, q: ArrayLike,
 def nanpercentile(a: ArrayLike, q: ArrayLike,
                   axis: int | tuple[int, ...] | None = None,
                   out: None = None, overwrite_input: bool = False, method: str = "linear",
-                  keepdims: bool = False, *, interpolation: str | DeprecatedArg = DeprecatedArg()) -> Array:
+                  keepdims: bool = False, weights: ArrayLike | None = None, *, interpolation: str | DeprecatedArg = DeprecatedArg()) -> Array:
   """Compute the percentile of the data along the specified axis, ignoring NaN values.
 
   JAX implementation of :func:`numpy.nanpercentile`.
@@ -2688,7 +2749,7 @@ def nanpercentile(a: ArrayLike, q: ArrayLike,
        "Use 'method=' instead."), stacklevel=2)
     method = interpolation
   return nanquantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
-                     method=method, keepdims=keepdims)
+                     method=method, keepdims=keepdims, weights=weights)
 
 
 @export

tests/lax_numpy_reducers_test.py

@@ -28,6 +28,7 @@
 
 from jax._src import config
 from jax._src import dtypes
+from jax._src.numpy.reductions import quantile
 from jax._src import test_util as jtu
 from jax._src.util import NumpyComplexWarning
 
@@ -763,6 +764,14 @@ def testPercentilePrecision(self):
     x = jnp.float64([1, 2, 3, 4, 7, 10])
     self.assertEqual(jnp.percentile(x, 50), 3.5)
 
+  def test_weighted_quantile_linear(self):
+    a = jnp.array([1, 2, 3, 4, 5], dtype=float)
+    weights = jnp.array([1, 2, 1, 1, 1], dtype=float)
+    q = jnp.array([0.5])
+    expected = np.quantile(a, q, weights=weights)
+    result = quantile(a, q, weights=weights, method="linear")
+    np.testing.assert_allclose(result, expected, rtol=1e-6)
+
   @jtu.sample_product(
     [dict(a_shape=a_shape, axis=axis)
       for a_shape, axis in (