{"record":{"id":"4c88b252f8aa72cb","repo":"jax-ml/jax","slug":"put-along-axis-argument-values-must-be-broadcast","errorCode":null,"errorMessage":"put_along_axis argument 'values' must be broadcastable to 'indices'. Got {values.shape=} and {indices.shape=}.","messagePattern":"put_along_axis argument 'values' must be broadcastable to 'indices'\\. Got (.+?) and (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":1057,"sourceCode":"  arr, indices, values = util.ensure_arraylike(\"put_along_axis\", arr, indices, values)\n\n  original_axis = axis\n  original_arr_shape = arr.shape\n\n  if axis is None:\n    arr = arr.ravel()\n    axis = 0\n\n  if not arr.ndim == indices.ndim:\n    raise ValueError(\n      \"put_along_axis arguments 'arr' and 'indices' must have same ndim. Got \"\n      f\"{arr.ndim=} and {indices.ndim=}.\"\n    )\n\n  try:\n    values = util._broadcast_to(values, indices.shape)\n  except ValueError:\n    raise ValueError(\n      \"put_along_axis argument 'values' must be broadcastable to 'indices'. Got \"\n      f\"{values.shape=} and {indices.shape=}.\"\n    )\n\n  idx = _make_along_axis_idx(arr.shape, indices, axis)\n  result = arr.at[idx].set(values, mode=mode)\n\n  if original_axis is None:\n    result = result.reshape(original_arr_shape)\n\n  return result\n\n\n### Indexing\n\ndef _is_integer_index(idx: Any) -> bool:\n  return isinstance(idx, (int, np.integer)) and not isinstance(idx, (bool, np.bool_))\n","sourceCodeStart":1039,"sourceCodeEnd":1075,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L1039-L1075","documentation":"values must broadcast to indices.shape for put_along_axis (updates land at index positions, so the values shape is governed by the index grid, not arr). This wraps the underlying broadcast error with shapes shown.","triggerScenarios":"jnp.put_along_axis(arr, indices, values, ...) where values.shape cannot broadcast to indices.shape — e.g. values (B,) with indices (B, N), or scalar-vs-grid mismatches are fine but (N,) vs (B, N) with N != B fails.","commonSituations":"Passing per-row values where a full grid is needed; reusing values shaped for arr rather than for the index grid.","solutions":["Reshape/broadcast values to indices.shape: values = jnp.broadcast_to(values, indices.shape)","Or use values[..., None] to add the trailing axis matching the index grid"],"exampleFix":"// before\na = jnp.put_along_axis(a, idx, vals, axis=1, inplace=False)  # vals (B,), idx (B,N)\n// after\na = jnp.put_along_axis(a, idx, jnp.broadcast_to(vals[:, None], idx.shape), axis=1, inplace=False)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nvalues_b = jnp.broadcast_to(values, indices.shape)  # raises early if incompatible","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pre-broadcast values to indices.shape before calling","Remember values follows the index grid, not arr's shape"],"tags":["jax","put-along-axis","broadcasting","shape-mismatch"],"backgroundTag":"broadcast-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}