Add weighted quantile and percentile support with tests

Author: Aniketsy

jax/_src/numpy/reductions.py

@@ -23,10 +23,12 @@
 
 import numpy as np
 
-import jax
-from jax import lax
-import jax._src.numpy as jnp
+from jax._src.lax import lax
+from jax._src.lax import parallel as lax_parallel
+from jax._src.lax import slicing
+from jax._src.lax.control_flow import loops
 from jax._src import api
+from jax._src import config
 from jax._src import core
 from jax._src import deprecations
 from jax._src import dtypes
@@ -63,7 +65,7 @@ def _moveaxis(a: ArrayLike, source: int, destination: int) -> Array:
   perm.insert(destination, source)
   return lax.transpose(a, perm)
 
-def _upcast_f16(dtype: DTypeLike) -> DType:
+def _upcast_f16(dtype: DTypeLike) -> DTypeLike:
   if np.dtype(dtype) in [np.float16, dtypes.bfloat16]:
     return np.dtype('float32')
   return np.dtype(dtype)
@@ -234,7 +236,7 @@ def _reduce_sum(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
   return _reduction(a, "sum", lax.add, 0, preproc=_cast_to_numeric,
                     bool_op=lax.bitwise_or, upcast_f16_for_computation=(dtype is None),
                     axis=axis, dtype=dtype, out=out, keepdims=keepdims,
-                    initial=initial, where_=where, parallel_reduce=lax.psum,
+                    initial=initial, where_=where, parallel_reduce=lax_parallel.psum,
                     promote_integers=promote_integers)
 
 
@@ -407,7 +409,7 @@ def _reduce_max(a: ArrayLike, axis: Axis = None, out: None = None,
                 where: ArrayLike | None = None) -> Array:
   return _reduction(a, "max", lax.max, -np.inf, has_identity=False,
                     axis=axis, out=out, keepdims=keepdims,
-                    initial=initial, where_=where, parallel_reduce=lax.pmax)
+                    initial=initial, where_=where, parallel_reduce=lax_parallel.pmax)
 
 
 @export
@@ -490,7 +492,7 @@ def _reduce_min(a: ArrayLike, axis: Axis = None, out: None = None,
                 where: ArrayLike | None = None) -> Array:
   return _reduction(a, "min", lax.min, np.inf, has_identity=False,
                     axis=axis, out=out, keepdims=keepdims,
-                    initial=initial, where_=where, parallel_reduce=lax.pmin)
+                    initial=initial, where_=where, parallel_reduce=lax_parallel.pmin)
 
 
 @export
@@ -797,7 +799,7 @@ def _axis_size(a: ArrayLike, axis: int | Sequence[int]):
   size = 1
   a_shape = np.shape(a)
   for a in axis_seq:
-    size *= maybe_named_axis(a, lambda i: a_shape[i], lambda name: lax.psum(1, name))
+    size *= maybe_named_axis(a, lambda i: a_shape[i], lambda name: lax_parallel.psum(1, name))
   return size
 
 
@@ -1140,12 +1142,12 @@ def _var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
   normalizer = lax.sub(normalizer, lax.convert_element_type(correction, computation_dtype))
   result = sum(centered, axis, dtype=computation_dtype, keepdims=keepdims, where=where)
   result = lax.div(result, normalizer).astype(dtype)
-  with jax.debug_nans(False):
+  with config.debug_nans(False):
     result = _where(normalizer > 0, result, np.nan)
   return result
 
 
-def _var_promote_types(a_dtype: DTypeLike, dtype: DTypeLike | None) -> tuple[DType, DType]:
+def _var_promote_types(a_dtype: DTypeLike, dtype: DTypeLike | None) -> tuple[Any, Any]:
   if dtype:
     if (not dtypes.issubdtype(dtype, np.complexfloating) and
         dtypes.issubdtype(a_dtype, np.complexfloating)):
@@ -2010,8 +2012,8 @@ def _cumulative_reduction(
   if fill_nan:
     a = _where(lax_internal._isnan(a), _lax_const(a, fill_value), a)
 
-  a_type: DType = dtypes.dtype(a)
-  result_type: DTypeLike = dtypes.dtype(dtype or a)
+  a_type = dtypes.dtype(a)
+  result_type = dtypes.dtype(dtype or a)
   if dtype is None and promote_integers or dtypes.issubdtype(result_type, np.bool_):
     result_type = _promote_integer_dtype(result_type)
   result_type = dtypes.canonicalize_dtype(result_type)
@@ -2062,7 +2064,7 @@ def cumsum(a: ArrayLike, axis: int | None = None,
     Array([[ 1,  3,  6],
            [ 4,  9, 15]], dtype=int32)
   """
-  return _cumulative_reduction("cumsum", lax.cumsum, a, axis, dtype, out)
+  return _cumulative_reduction("cumsum", loops.cumsum, a, axis, dtype, out)
 
 
 @export
@@ -2098,7 +2100,7 @@ def cumprod(a: ArrayLike, axis: int | None = None,
     Array([[  1,   2,   6],
            [  4,  20, 120]], dtype=int32)
   """
-  return _cumulative_reduction("cumprod", lax.cumprod, a, axis, dtype, out)
+  return _cumulative_reduction("cumprod", loops.cumprod, a, axis, dtype, out)
 
 
 @export
@@ -2147,7 +2149,7 @@ def nancumsum(a: ArrayLike, axis: int | None = None,
     Array([[ 1.,  3.,  3.],
            [ 4.,  4., 10.]], dtype=float32)
   """
-  return _cumulative_reduction("nancumsum", lax.cumsum, a, axis, dtype, out,
+  return _cumulative_reduction("nancumsum", loops.cumsum, a, axis, dtype, out,
                                fill_nan=True, fill_value=0)
 
 
@@ -2196,15 +2198,15 @@ def nancumprod(a: ArrayLike, axis: int | None = None,
     Array([[ 1.,  2.,  2.],
            [ 4.,  4., 24.]], dtype=float32)
   """
-  return _cumulative_reduction("nancumprod", lax.cumprod, a, axis, dtype, out,
+  return _cumulative_reduction("nancumprod", loops.cumprod, a, axis, dtype, out,
                                fill_nan=True, fill_value=1)
 
 
 @partial(api.jit, static_argnames=('axis', 'dtype'))
 def _cumsum_with_promotion(a: ArrayLike, axis: int | None = None,
            dtype: DTypeLike | None = None, out: None = None) -> Array:
   """Utility function to compute cumsum with integer promotion."""
-  return _cumulative_reduction("_cumsum_with_promotion", lax.cumsum,
+  return _cumulative_reduction("_cumsum_with_promotion", loops.cumsum,
                                a, axis, dtype, out, promote_integers=True)
 
 
@@ -2322,7 +2324,7 @@ def cumulative_prod(
 
   axis = _canonicalize_axis(axis, x.ndim)
   dtypes.check_user_dtype_supported(dtype)
-  out = _cumulative_reduction("cumulative_prod", lax.cumprod, x, axis, dtype)
+  out = _cumulative_reduction("cumulative_prod", loops.cumprod, x, axis, dtype)
   if include_initial:
     zeros_shape = list(x.shape)
     zeros_shape[axis] = 1
@@ -2486,21 +2488,24 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
 
   q, = promote_dtypes_inexact(q)
   q = lax_internal.asarray(q)
-  if getattr(q, "ndim", 0) == 0:
+  q_was_scalar = (getattr(q, "ndim", 0) == 0)
+  if q_was_scalar:
     q = lax.expand_dims(q, (0,))
   q_shape = q.shape
   q_ndim = q.ndim
   if q_ndim > 1:
     raise ValueError(f"q must be have rank <= 1, got shape {q.shape}")
-
   a_shape = a.shape
   # Handle weights
   if weights is None:
     a, = promote_dtypes_inexact(a)
   else:
-    a, weights = promote_dtypes_inexact(a, weights)
-    weights = lax.convert_element_type(weights, a.dtype)
-    a_shape = a.shape
+    common_dtype = np.result_type(a, q, weights, np.float32)
+    a = a.astype(common_dtype)
+    q = q.astype(common_dtype)
+    weights = weights.astype(common_dtype)
+    a,q, weights = promote_dtypes_inexact(a, q, weights)
+    #weights = lax.convert_element_type(weights, a.dtype)
     w_shape = np.shape(weights)
     if np.ndim(weights) == 0:
       weights = lax.broadcast_in_dim(weights, a_shape, ())
@@ -2511,8 +2516,9 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
       if axis is None:
         raise TypeError("Axis must be specified when shapes of a and weights differ.")
       if isinstance(axis, tuple):
-        if w_shape != tuple(a_shape[i] for i in axis):
-            raise ValueError("Shape of weights must match the shape of the axes being reduced.")
+        expected_shape = tuple(a_shape[i] for i in axis)
+        if w_shape != expected_shape:
+          raise ValueError("Shape of weights must match the shape of the axes being reduced.")
         weights = lax.broadcast_in_dim(
             weights,
             shape=a_shape,
@@ -2521,18 +2527,23 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
         w_shape = a_shape
       else:
         if len(w_shape) != 1 or w_shape[0] != a_shape[axis]:
-            raise ValueError("Length of weights not compatible with specified axis.")
+          raise ValueError("Length of weights not compatible with specified axis.")
         weights = lax.expand_dims(weights, (axis,))
         weights = _broadcast_to(weights, a.shape)
         w_shape = a_shape
-
+    
     if squash_nans:
       nan_mask = ~lax_internal._isnan(a)
       weights = _where(nan_mask, weights, 0)
     else:
-      with jax.debug_nans(False):
+      with config.debug_nans(False):
         a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
 
+    if all(weights == 0):
+      raise ValueError("Sum of weights must not be zero")
+    if any(weights < 0):
+      raise ValueError("Weights must be non-negative")
+
     total_weight = sum(weights, axis=axis, keepdims=True)
     a_sorted, weights_sorted = lax.sort_key_val(a, weights, dimension=axis)
     cum_weights = cumsum(weights_sorted, axis=axis)
@@ -2549,15 +2560,15 @@ def _weighted_quantile(qi):
       slice_sizes[axis] = 1
       offset_start = q_ndim
       total_offset_dims = len(a_shape) + q_ndim if keepdims else len(a_shape) + q_ndim - 1
-      dnums = lax.GatherDimensionNumbers(
+      dnums = slicing.GatherDimensionNumbers(
         offset_dims=tuple(range(offset_start, total_offset_dims)),
         collapsed_slice_dims=(axis,),
         start_index_map=(axis,)
      )
-      val = lax.gather(a_sorted, idx[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
-      val_prev = lax.gather(a_sorted, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
-      cw_prev = lax.gather(cum_weights_norm, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
-      cw_next = lax.gather(cum_weights_norm, idx[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
+      val = slicing.gather(a_sorted, idx[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
+      val_prev = slicing.gather(a_sorted, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
+      cw_prev = slicing.gather(cum_weights_norm, idx_prev[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
+      cw_next = slicing.gather(cum_weights_norm, idx[..., None], dimension_numbers=dnums, slice_sizes=slice_sizes)
       if method == "linear":
         denom = cw_next - cw_prev
         denom = _where(denom == 0, 1, denom)
@@ -2577,16 +2588,15 @@ def _weighted_quantile(qi):
         raise ValueError(f"{method=!r} not recognized")
       return out
 
-    result = jax.vmap(_weighted_quantile)(q)
+    result = api.vmap(_weighted_quantile)(q)
     if keepdims and keepdim:
-      if q_ndim > 0:
-        keepdim = [q_shape[0], *keepdim]
-      result = result.reshape(tuple(keepdim))
-    else:
-      if q_ndim == 0 or (q_ndim == 1 and q_shape[0] == 1):
-        if result.ndim > 0 and result.shape[0] == 1:
-          result = lax.squeeze(result, (0,))
-    return lax.convert_element_type(result, a.dtype)
+      keepdim_out = list(keepdim)
+      if not q_was_scalar:
+        keepdim_out = [q_shape[0], *keepdim_out]
+      result = result.reshape(tuple(keepdim_out))
+    elif q_was_scalar and result.ndim > 0 and result.shape[0] == 1:
+      result = result.squeeze(axis=0)
+    return result
 
   if squash_nans:
     a = _where(lax_internal._isnan(a), np.nan, a) # Ensure nans are positive so they sort to the end.
@@ -2617,7 +2627,7 @@ def _weighted_quantile(qi):
     index[axis] = high
     high_value = a[tuple(index)]
   else:
-    with jax.debug_nans(False):
+    with config.debug_nans(False):
       a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
     a = lax.sort(a, dimension=axis)
     n = lax.convert_element_type(a_shape[axis], lax_internal._dtype(q))
@@ -2634,15 +2644,15 @@ def _weighted_quantile(qi):
 
     slice_sizes = list(a_shape)
     slice_sizes[axis] = 1
-    dnums = lax.GatherDimensionNumbers(
+    dnums = slicing.GatherDimensionNumbers(
       offset_dims=tuple(range(
         q_ndim,
         len(a_shape) + q_ndim if keepdims else len(a_shape) + q_ndim - 1)),
       collapsed_slice_dims=() if keepdims else (axis,),
       start_index_map=(axis,))
-    low_value = lax.gather(a, low[..., None], dimension_numbers=dnums,
+    low_value = slicing.gather(a, low[..., None], dimension_numbers=dnums,
                            slice_sizes=slice_sizes)
-    high_value = lax.gather(a, high[..., None], dimension_numbers=dnums,
+    high_value = slicing.gather(a, high[..., None], dimension_numbers=dnums,
                             slice_sizes=slice_sizes)
     if q_ndim == 1:
       low_weight = lax.broadcast_in_dim(low_weight, low_value.shape,
@@ -2667,14 +2677,13 @@ def _weighted_quantile(qi):
   else:
     raise ValueError(f"{method=!r} not recognized")
   if keepdims and keepdim:
-    if q_ndim > 0:
-      keepdim = [np.shape(q)[0], *keepdim]
-    result = result.reshape(keepdim)
-  else:
-    if q_ndim == 0 or (q_ndim == 1 and q_shape[0] == 1):
-      if result.ndim > 0 and result.shape[0] == 1:
-        result = lax.squeeze(result, (0,))
-  return lax.convert_element_type(result, a.dtype)
+    keepdim_out = list(keepdim)
+    if not q_was_scalar:
+      keepdim_out = [q_shape[0], *keepdim_out]
+    result = result.reshape(tuple(keepdim_out))
+  elif q_was_scalar and result.ndim > 0 and result.shape[0] == 1:
+    result = result.squeeze(axis=0)
+  return result
 
 
 # TODO(jakevdp): interpolation argument deprecated 2024-05-16

tests/lax_numpy_reducers_test.py

@@ -795,7 +795,7 @@ def testWeightedQuantile(self, a_shape, a_dtype, q_shape, q_dtype, axis, keepdim
     weights = np.abs(rng(weights_shape, a_dtype)) + 1e-3
 
     def np_fun(a, q, weights):
-      return np.quantile(np.array(a), np.array(q), axis=axis, weights=np.array(weights), method=method, keepdims=keepdims)
+      return np.quantile(np.array(a), np.array(q), axis=axis, weights=weights, method=method, keepdims=keepdims)
     def jnp_fun(a, q, weights):
       return jnp.quantile(a, q, axis=axis, weights=weights, method=method, keepdims=keepdims)
     args_maker = lambda: [