jax-ml/jax · error · NotImplementedError

{method} doesn't support weights. Only method 'inverted_cdf'

Error message

{method} doesn't support weights. Only method 'inverted_cdf' supports weights.

What it means

Weighted quantiles in JAX are only implemented for method='inverted_cdf'. Requesting weights with any other method (the default 'linear' included) raises NotImplementedError.

Source

Thrown at jax/_src/numpy/reductions.py:2523

  """
  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(a, 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: 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:
    if dtypes.issubdtype(weights.dtype, np.complexfloating):
      raise ValueError("Weights cannot be complex types.")
    if method != "inverted_cdf":
      raise NotImplementedError(f"{method} doesn't support weights. Only method 'inverted_cdf' supports weights.")
    a, weights = promote_dtypes_inexact(a, weights)
    if weights.shape != a.shape:
      if axis is None:
        raise ValueError("Weights shape must match 'a' shape when axis is None.")
      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)}")
      weights = lax.broadcast_in_dim(weights, a.shape, broadcast_dimensions=ax_tuple)
  else:
    a, = promote_dtypes_inexact(a)
  keepdim = []
  if dtypes.issubdtype(a.dtype, np.complexfloating):
    raise ValueError("quantile does not support complex input, as the operation is poorly defined.")
  if axis is None:
    if keepdims:
      keepdim = [1] * a.ndim
    a = a.ravel()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add method='inverted_cdf' to the call
2. If interpolated weighted quantiles are required, implement manually (e.g. weighted cumulative distribution + interpolation) or compute on host with NumPy

Example fix

// before
jnp.quantile(a, q, weights=w)
// after
jnp.quantile(a, q, weights=w, method='inverted_cdf')

if weights is not None:
    method = 'inverted_cdf'
jnp.quantile(a, q, weights=weights, method=method)

When it happens

Trigger: Calling jnp.quantile(a, q, weights=w) without setting method (defaults to 'linear'), or with method='lower'/'higher'/'midpoint'/'nearest'.

Common situations: Assuming NumPy-style weighted quantiles work with default interpolation; enabling weights in an existing quantile call during feature work.

Related errors

• Weights cannot be complex types.
• Weights shape must match 'a' shape when axis is None.
• Weights shape {weights.shape} must match reduction axes {tup
• expected a 1-d array for weights
• expected w and y to have the same length

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/33c54adf6a8d0aeb. Report an issue: GitHub.