{"record":{"id":"fa7dba06b25fb028","repo":"jax-ml/jax","slug":"axis-must-be-specified-when-shapes-of-a-and-weight","errorCode":null,"errorMessage":"Axis must be specified when shapes of a and weights differ.","messagePattern":"Axis must be specified when shapes of a and weights differ\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/reductions.py","lineNumber":1017,"sourceCode":"def _average(a: ArrayLike, axis: Axis = None, weights: ArrayLike | None = None,\n             returned: bool = False, keepdims: bool = False) -> Array | tuple[Array, Array]:\n  axis_tuple = canonicalize_axis_tuple(axis, np.ndim(a))\n\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","sourceCodeStart":999,"sourceCodeEnd":1035,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/reductions.py#L999-L1035","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Specify the axis the weights apply to: jnp.average(a, axis=0, weights=w)","Broadcast weights to a.shape explicitly if they apply elementwise","If you want plain weighted mean over all elements, flatten both: jnp.average(a.ravel(), weights=w.ravel())"],"exampleFix":"// before\njnp.average(a, weights=w)  # a (3,4), w (3,)\n// after\njnp.average(a, axis=0, weights=w)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\na, w = jnp.asarray(a), jnp.asarray(w)\nif a.shape != w.shape and axis is None:\n    axis = int(jnp.argmin([abs(a.ndim - 1), 1]))  # or explicitly pick the weighted axis\njnp.average(a, axis=axis, weights=w)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass axis when weights are 1-D","Broadcast weights manually if elementwise"],"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"}