jax-ml/jax · error · ValueError

Weights shape {weights.shape} must match reduction axes {tup

Error message

Weights shape {weights.shape} must match reduction axes {tuple(a.shape[ax] for ax in ax_tuple)}

What it means

When an explicit axis is given to a weighted quantile, weights must have shape equal to the reduction axes of a (e.g. a.shape[axis] for a single axis). Otherwise JAX raises ValueError showing both shapes.

Source

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

    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):
    keepdim = list(a.shape)
    nd = a.ndim
    axis = tuple(canonicalize_axis(ax, nd) for ax in axis)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape weights to the reduction axes: w = w.reshape(a.shape[axis]) for a single axis
2. Or keep weights full-shaped and drop axis (axis=None) so they match a.shape
3. Double-check axis orientation vs weight layout with a shape assertion before calling

Example fix

// before
jnp.quantile(a, q, axis=0, weights=w, method='inverted_cdf')  # w has shape of full a
// after
jnp.quantile(a, q, axis=0, weights=w.reshape(a.shape[0]), method='inverted_cdf')

import jax.numpy as jnp

ax = jnp.canonicalize_axis(axis, a.ndim)
if weights.shape != (a.shape[ax],):
    weights = weights.reshape(a.shape[ax])
jnp.quantile(a, q, axis=axis, weights=weights, method='inverted_cdf')

def weights_match_reduction_axes(w, a, axis) -> bool:
    ax = jnp.canonicalize_axis(axis, a.ndim)
    return w.shape == (a.shape[ax],)

When it happens

Trigger: Calling jnp.quantile(a, q, axis=0, weights=w, method='inverted_cdf') where w.shape != (a.shape[0],), e.g. full a-shaped weights with axis set, or transposed weights.

Common situations: Switching a working axis=None call to a per-axis reduction without reshaping weights; transposition bugs where weights align to the wrong dimension.

Related errors

• Weights shape must match 'a' shape when axis is None.
• 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/19926a24faaa8950. Report an issue: GitHub.