Address review comments

Author: C-you-know

jax/_src/numpy/reductions.py

@@ -2375,8 +2375,8 @@ def cumulative_prod(
 @export
 @api.jit(static_argnames=('axis', 'overwrite_input', 'keepdims', 'method'))
 def quantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = None,
-             weights: ArrayLike | None = None, out: None = None, overwrite_input: bool = False, method: str = "linear",
-             keepdims: bool = False) -> Array:
+             out: None = None, overwrite_input: bool = False, method: str = "linear",
+             keepdims: bool = False, *, weights: ArrayLike | None = None) -> Array:
   """Compute the quantile of the data along the specified axis.
 
   JAX implementation of :func:`numpy.quantile`.
@@ -2426,6 +2426,8 @@ def quantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = No
     Array(3., dtype=float32)
   """
   a, q = ensure_arraylike("quantile", a, q)
+  if weights is not None:
+    weights = ensure_arraylike("quantile", weights) 
   if overwrite_input or out is not None:
     raise ValueError("jax.numpy.quantile does not support overwrite_input=True "
                      "or out != None")
@@ -2435,8 +2437,8 @@ def quantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = No
 @export
 @api.jit(static_argnames=('axis', 'overwrite_input', 'keepdims', 'method'))
 def nanquantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = None,
-                weights: ArrayLike | None = None, out: None = None, overwrite_input: bool = False, method: str = "linear",
-                keepdims: bool = False) -> Array:
+                out: None = None, overwrite_input: bool = False, method: str = "linear",
+                keepdims: bool = False, *, weights: ArrayLike | None = None) -> Array:
   """Compute the quantile of the data along the specified axis, ignoring NaNs.
 
   JAX implementation of :func:`numpy.nanquantile`.
@@ -2487,18 +2489,19 @@ def nanquantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None =
     Array(4.0, dtype=float32)
   """
   a, q = ensure_arraylike("nanquantile", a, q)
+  if weights is not None:
+    weights = ensure_arraylike("nanquantile", weights)
   if overwrite_input or out is not None:
     msg = ("jax.numpy.nanquantile does not support overwrite_input=True or "
            "out != None")
     raise ValueError(msg)
   return _quantile(lax.asarray(a), lax.asarray(q), axis, method, keepdims, True, weights)
 
 def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
-              method: str, keepdims: bool, squash_nans: bool, weights: ArrayLike | None = None) -> Array:
+              method: str, keepdims: bool, squash_nans: bool, weights: Array | None = None) -> Array:
   if method not in ["linear", "lower", "higher", "midpoint", "nearest", "inverted_cdf"]:
     raise ValueError("method can only be 'linear', 'lower', 'higher', 'midpoint', 'nearest' or 'inverted_cdf'")
   if weights is not None:
-    weights = ensure_arraylike("_quantile", weights)
     weights = lax.asarray(weights)
     if dtypes.issubdtype(weights.dtype, np.complexfloating):
       raise ValueError("Weights cannot be complex types.")
@@ -2508,8 +2511,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     if weights.shape != a.shape:
       if axis is None:
         raise ValueError("Weights shape must match 'a' shape when axis is None.")
-      ax_tuple = (axis,) if isinstance(axis, int) else tuple(axis)
-      ax_tuple = tuple(canonicalize_axis(ax, a.ndim) for ax in ax_tuple)
+      ax_tuple = canonicalize_axis_tuple(axis, a.ndim)
       if weights.shape != tuple(a.shape[ax] for ax in ax_tuple):
         raise ValueError(f"Weights shape {weights.shape} must match reduction axes "
                           f"{tuple(a.shape[ax] for ax in ax_tuple)}")
@@ -2524,7 +2526,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
       keepdim = [1] * a.ndim
     a = a.ravel()
     if weights is not None:
-        weights = weights.ravel()
+      weights = weights.ravel()
     axis = 0
   elif isinstance(axis, tuple):
     keepdim = list(a.shape)
@@ -2544,7 +2546,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     touch_shape = tuple(x for idx,x in enumerate(a.shape) if idx in axis)
     a = lax.reshape(a, do_not_touch_shape + (math.prod(touch_shape),), dimensions)
     if weights is not None:
-        weights = lax.reshape(weights, do_not_touch_shape + (math.prod(touch_shape),), dimensions)
+      weights = lax.reshape(weights, do_not_touch_shape + (math.prod(touch_shape),), dimensions)
     axis = canonicalize_axis(-1, a.ndim)
   else:
     axis = canonicalize_axis(axis, a.ndim)
@@ -2559,9 +2561,9 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
   if squash_nans:
     a = _where(lax._isnan(a), np.nan, a) # Ensure nans are positive so they sort to the end.
     if weights is not None:
-        a, weights = lax.sort_key_val(a, weights, dimension=axis)
+      a, weights = lax.sort_key_val(a, weights, dimension=axis)
     else:
-        a = lax.sort(a, dimension=axis)
+      a = lax.sort(a, dimension=axis)
     counts = sum(lax.bitwise_not(lax._isnan(a)), axis=axis, dtype=q.dtype, keepdims=keepdims)
     shape_after_reduction = counts.shape
     q = lax.expand_dims(
@@ -2591,9 +2593,9 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     with config.debug_nans(False):
       a = _where(any(lax._isnan(a), axis=axis, keepdims=True), np.nan, a)
     if weights is not None:
-        a, weights = lax.sort_key_val(a, weights, dimension=axis)
+      a, weights = lax.sort_key_val(a, weights, dimension=axis)
     else:
-        a = lax.sort(a, dimension=axis)
+      a = lax.sort(a, dimension=axis)
     n = lax.convert_element_type(a_shape[axis], lax._dtype(q))
     q = lax.mul(q, n - 1)
     low = lax.floor(q)
@@ -2638,7 +2640,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     result = lax.mul(lax.add(low_value, high_value), lax._const(low_value, 0.5))
   elif method == "inverted_cdf":
     if weights is None:
-      weights = lax.full(a.shape, 1.0, dtype=a.dtype)
+      weights = lax.full_like(a, 1.0)
     zeros = lax.full_like(weights, 0)
     bad_weights = lax.bitwise_or(lax.lt(weights, zeros), lax._isnan(weights))
     nan_data = lax._isnan(a)
@@ -2651,13 +2653,13 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     target_w_aligned = target_w if keepdims else lax.expand_dims(target_w, (axis + q_ndim,))
     cw_f = lax.expand_dims(cum_weights, tuple(range(q_ndim)))
     is_less = lax.lt(cw_f, target_w_aligned)
-    idx = sum(lax.convert_element_type(is_less, np.int32), axis=axis + q_ndim, keepdims=keepdims)
+    idx = sum(lax.convert_element_type(is_less, dtypes.default_int_dtype()), axis=axis + q_ndim, keepdims=keepdims)
     if squash_nans:
       valid_counts = sum(lax.bitwise_not(nan_data), axis=axis, dtype=q.dtype, keepdims=keepdims)
     else:
       valid_counts = lax.full_like(total_weight, a_shape[axis], dtype=q.dtype)
     limit = lax.sub(valid_counts, lax._const(valid_counts, 1))
-    max_idx = lax.convert_element_type(limit, np.int32)
+    max_idx = lax.convert_element_type(limit, dtypes.default_int_dtype())
     max_idx_f = lax.expand_dims(max_idx, tuple(range(q_ndim)))
     max_idx_f = lax.convert_element_type(max_idx_f, idx.dtype)
     idx = lax.max(lax._const(idx, 0), lax.min(idx, max_idx_f))
@@ -2670,9 +2672,9 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
       result = indexing.take_along_axis(a, idx_transposed, axis=axis)
       result = lax.squeeze(result, (axis,))
     else:
-      perm = (list(range(q_ndim, q_ndim + axis)) +
-              list(range(q_ndim)) +
-              list(range(q_ndim + axis, idx_take.ndim)))
+      perm = [*range(q_ndim, q_ndim + axis),
+              *range(q_ndim),
+              *range(q_ndim + axis, idx_take.ndim)]
       idx_transposed = lax.transpose(idx_take, perm)
       result = indexing.take_along_axis(a, idx_transposed, axis=axis)
       inv_perm = [perm.index(i) for i in range(len(perm))]
@@ -2701,7 +2703,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",
-               weights: ArrayLike | None = None, keepdims: bool = False) -> Array:
+               keepdims: bool = False, *, weights: ArrayLike | None = None) -> Array:
   """Compute the percentile of the data along the specified axis.
 
   JAX implementation of :func:`numpy.percentile`.
@@ -2751,6 +2753,8 @@ def percentile(a: ArrayLike, q: ArrayLike,
     Array(3., dtype=float32)
   """
   a, q = ensure_arraylike("percentile", a, q)
+  if weights is not None:
+    weights = ensure_arraylike("percentile", weights)
   q, = promote_dtypes_inexact(q)
   return quantile(a, q / 100, axis=axis, out=out, overwrite_input=overwrite_input,
                   method=method, weights=weights, keepdims=keepdims)
@@ -2761,7 +2765,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",
-                  weights: ArrayLike | None = None, keepdims: bool = False) -> Array:
+                  keepdims: bool = False, *, weights: ArrayLike | None = None) -> Array:
   """Compute the percentile of the data along the specified axis, ignoring NaN values.
 
   JAX implementation of :func:`numpy.nanpercentile`.
@@ -2813,6 +2817,8 @@ def nanpercentile(a: ArrayLike, q: ArrayLike,
     Array(4.0, dtype=float32)
   """
   a, q = ensure_arraylike("nanpercentile", a, q)
+  if weights is not None:
+    weights = ensure_arraylike("nanpercentile", weights)
   q, = promote_dtypes_inexact(q)
   q = q / 100
   return nanquantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,