{"record":{"id":"428004b2dca90c67","repo":"jax-ml/jax","slug":"shape-of-weights-must-be-consistent-with-shape-of","errorCode":null,"errorMessage":"Shape of weights must be consistent with shape of a along specified axis.","messagePattern":"Shape of weights must be consistent with shape of a along specified axis\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/reductions.py","lineNumber":1020,"sourceCode":"\n  if weights is None: # Treat all weights as 1\n    a = ensure_arraylike(\"average\", a)\n    a, = promote_dtypes_inexact(a)\n    avg = mean(a, axis=axis, keepdims=keepdims)\n    if axis is None:\n      weights_sum = lax.full((), core.dimension_as_value(a.size), dtype=avg.dtype)\n    else:\n      weights_sum = lax.full((), math.prod(core.dimension_as_value(a.shape[d]) for d in axis_tuple), dtype=avg.dtype)\n  else:\n    a, weights = ensure_arraylike(\"average\", a, weights)\n    a, weights = promote_dtypes_inexact(a, weights)\n\n    if a.shape != weights.shape:\n      if axis is None:\n        raise ValueError(\"Axis must be specified when shapes of a and \"\n                         \"weights differ.\")\n      if weights.shape != tuple(a.shape[ax] for ax in axis_tuple):\n        raise ValueError(\"Shape of weights must be consistent with shape \"\n                         \"of a along specified axis.\")\n      new_shape = tuple(dim if i in axis_tuple else 1 for i, dim in enumerate(a.shape))\n      weights = lax.reshape(weights, new_shape, dimensions=tuple(np.argsort(axis_tuple)))\n\n    weights_sum = sum(weights, axis=axis, keepdims=keepdims)\n    avg = sum(a * weights, axis=axis, keepdims=keepdims) / weights_sum\n\n  if returned:\n    if avg.shape != weights_sum.shape:\n      weights_sum = _broadcast_to(weights_sum, avg.shape)\n    return avg, weights_sum\n  return avg\n\n\n@export\ndef var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,\n        out: None = None, ddof: int = 0, keepdims: bool = False, *,\n        where: ArrayLike | None = None, mean: ArrayLike | None = None,","sourceCodeStart":1002,"sourceCodeEnd":1038,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/reductions.py#L1002-L1038","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Fix the axis or the weights so weights.shape == a.shape[axis]","Use w.shape[0] == a.shape[axis] check before calling","If weighting multiple axes, pass full-shape weights (no axis needed)"],"exampleFix":"// before\njnp.average(a, axis=0, weights=w)  # a (3,4), w has length 4\n// after\njnp.average(a, axis=0, weights=w)  # ensure len(w) == a.shape[0] == 3","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\naxis_t = (axis,) if isinstance(axis, int) else tuple(axis)\nexpected = tuple(jnp.asarray(a).shape[ax] for ax in axis_t)\nassert jnp.asarray(weights).shape == expected, f'weights {weights.shape} != {expected}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Double-check weights length equals a.shape[axis]","Verify orientation after transposing data"],"tags":["jax","average","weights","shape-mismatch"],"backgroundTag":"weights-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}