jax-ml/jax · error · ValueError

Axis must be specified when shapes of a and weights differ.

Error message

Axis must be specified when shapes of a and weights differ.

What it means

jnp.average requires the axis argument when a and weights have different shapes: without an axis, JAX cannot determine along which dimensions the weights apply, since only full-shape matching weights are unambiguous.

Source

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

def _average(a: ArrayLike, axis: Axis = None, weights: ArrayLike | None = None,
             returned: bool = False, keepdims: bool = False) -> Array | tuple[Array, Array]:
  axis_tuple = canonicalize_axis_tuple(axis, np.ndim(a))

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Specify the axis the weights apply to: jnp.average(a, axis=0, weights=w)
2. Broadcast weights to a.shape explicitly if they apply elementwise
3. If you want plain weighted mean over all elements, flatten both: jnp.average(a.ravel(), weights=w.ravel())

Example fix

// before
jnp.average(a, weights=w)  # a (3,4), w (3,)
// after
jnp.average(a, axis=0, weights=w)

import jax.numpy as jnp
a, w = jnp.asarray(a), jnp.asarray(w)
if a.shape != w.shape and axis is None:
    axis = int(jnp.argmin([abs(a.ndim - 1), 1]))  # or explicitly pick the weighted axis
jnp.average(a, axis=axis, weights=w)

When it happens

Trigger: jnp.average(a, weights=w) where a.shape != w.shape and axis=None, e.g. averaging a (3, 4) array with 1-D weights of length 3.

Common situations: Row/column weighting of 2-D data (most common): weights for one axis but axis omitted; porting numpy where the same error occurs; weights computed from a different reduction axis than intended.

Related errors

• Shape of weights must be consistent with shape of a along sp
• 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/fa7dba06b25fb028. Report an issue: GitHub.