jax-ml/jax · error · ValueError

Weights shape must match 'a' shape when axis is None.

Error message

Weights shape must match 'a' shape when axis is None.

What it means

When axis=None, weighted quantile requires weights with exactly the same shape as a (the reduction is over the whole array). A shape mismatch raises ValueError.

Source

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

  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()
    if weights is not None:
      weights = weights.ravel()
    axis = 0
  elif isinstance(axis, tuple):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast weights to a.shape first: w = jnp.broadcast_to(w, a.shape)
2. Pass an explicit axis and supply weights matching just the reduction axes
3. Reshape weights: w.reshape(a.shape) when sizes match element-wise

Example fix

// before
jnp.quantile(a, q, weights=w, method='inverted_cdf')  # w.shape != a.shape, axis=None
// after
jnp.quantile(a, q, weights=jnp.broadcast_to(w, a.shape), method='inverted_cdf')

import jax.numpy as jnp

if axis is None and weights is not None and weights.shape != a.shape:
    weights = jnp.broadcast_to(weights, a.shape)
jnp.quantile(a, q, axis=axis, weights=weights, method='inverted_cdf')

def weights_match_full_shape(w, a) -> bool:
    return w.shape == a.shape

try:
    jnp.quantile(a, q, weights=w, method='inverted_cdf')
except ValueError as e:
    if 'Weights shape' in str(e):
        w = jnp.broadcast_to(w, a.shape)
        q_val = jnp.quantile(a, q, weights=w, method='inverted_cdf')
    else:
        raise

When it happens

Trigger: Calling jnp.quantile(a, q, axis=None, weights=w, method='inverted_cdf') where w.shape != a.shape, e.g. flat weights against a 2-d array.

Common situations: Passing per-feature weight vectors to a whole-array quantile; reshaping a for a batched pipeline while keeping old 1-d weights.

Related errors

• Weights shape {weights.shape} must match reduction axes {tup
• Weights cannot be complex types.
• {method} doesn't support weights. Only method 'inverted_cdf'
• type of weights must match type of x. Got typeof(x)={core.ty
• multi_dot: last dimension of each array must match first dim

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