jax-ml/jax · error · ValueError

Shape of weights must be consistent with shape of a along sp

Error message

Shape of weights must be consistent with shape of a along specified axis.

What it means

When shapes of a and weights differ, the weights must match the size of a along the specified axis (weights.shape == tuple(a.shape[ax] for ax in axis)); otherwise the weighting is ill-defined.

Source

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

  if weights is None: # Treat all weights as 1
    a = ensure_arraylike("average", a)
    a, = promote_dtypes_inexact(a)
    avg = mean(a, axis=axis, keepdims=keepdims)
    if axis is None:
      weights_sum = lax.full((), core.dimension_as_value(a.size), dtype=avg.dtype)
    else:
      weights_sum = lax.full((), math.prod(core.dimension_as_value(a.shape[d]) for d in axis_tuple), dtype=avg.dtype)
  else:
    a, weights = ensure_arraylike("average", a, weights)
    a, weights = promote_dtypes_inexact(a, weights)

    if a.shape != weights.shape:
      if axis is None:
        raise ValueError("Axis must be specified when shapes of a and "
                         "weights differ.")
      if weights.shape != tuple(a.shape[ax] for ax in axis_tuple):
        raise ValueError("Shape of weights must be consistent with shape "
                         "of a along specified axis.")
      new_shape = tuple(dim if i in axis_tuple else 1 for i, dim in enumerate(a.shape))
      weights = lax.reshape(weights, new_shape, dimensions=tuple(np.argsort(axis_tuple)))

    weights_sum = sum(weights, axis=axis, keepdims=keepdims)
    avg = sum(a * weights, axis=axis, keepdims=keepdims) / weights_sum

  if returned:
    if avg.shape != weights_sum.shape:
      weights_sum = _broadcast_to(weights_sum, avg.shape)
    return avg, weights_sum
  return avg


@export
def var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
        out: None = None, ddof: int = 0, keepdims: bool = False, *,
        where: ArrayLike | None = None, mean: ArrayLike | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the axis or the weights so weights.shape == a.shape[axis]
2. Use w.shape[0] == a.shape[axis] check before calling
3. If weighting multiple axes, pass full-shape weights (no axis needed)

Example fix

// before
jnp.average(a, axis=0, weights=w)  # a (3,4), w has length 4
// after
jnp.average(a, axis=0, weights=w)  # ensure len(w) == a.shape[0] == 3

import jax.numpy as jnp
axis_t = (axis,) if isinstance(axis, int) else tuple(axis)
expected = tuple(jnp.asarray(a).shape[ax] for ax in axis_t)
assert jnp.asarray(weights).shape == expected, f'weights {weights.shape} != {expected}'

When it happens

Trigger: jnp.average(a, axis=0, weights=w) with a shape (3, 4) but w of length 4 (should be 3 for axis 0); multi-axis averaging where weights match only one of the axes.

Common situations: Transposed data (weights computed for rows but axis=0 iterates rows vs columns confusion); off-by-one in weight vectors; multi-axis tuples where weights match only a sub-axis.

Related errors

• Axis must be specified when shapes of a and weights differ.
• type of weights must match type of x. Got typeof(x)={core.ty
• Weights shape must match 'a' shape when axis is None.
• Weights shape {weights.shape} must match reduction axes {tup
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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