Add weighted quantile and percentile support with tests

Author: Aniketsy

jax/_src/numpy/reductions.py

@@ -2497,8 +2497,8 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
   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, q, weights = promote_dtypes_inexact(a, q, weights)
+    #weights = lax.convert_element_type(weights, a.dtype)
     a_shape = a.shape
     w_shape = np.shape(weights)
     if np.ndim(weights) == 0:
@@ -2583,8 +2583,6 @@ def _weighted_quantile(qi):
     elif q_was_scalar and result.ndim > 0 and result.shape[0] == 1:
       result = result.squeeze(axis=0)
     return result
-    return result
-
 
   if squash_nans:
     a = _where(lax_internal._isnan(a), np.nan, a) # Ensure nans are positive so they sort to the end.

tests/lax_numpy_reducers_test.py

@@ -789,7 +789,7 @@ def testWeightedQuantile(self, a_shape, a_dtype, q_shape, q_dtype, axis, keepdim
     if axis is None:
       weights_shape = a_shape
     elif isinstance(axis, tuple):
-      weights_shape = tuple(a_shape[i] for i in axis)
+      weights_shape = a_shape
     else:
       weights_shape = (a_shape[axis],)
     weights = np.abs(rng(weights_shape, a_dtype)) + 1e-3