{"record":{"id":"8396e04eea6c56fb","repo":"jax-ml/jax","slug":"put-along-axis-arguments-arr-and-indices-must","errorCode":null,"errorMessage":"put_along_axis arguments 'arr' and 'indices' must have same ndim. Got {arr.ndim=} and {indices.ndim=}.","messagePattern":"put_along_axis arguments 'arr' and 'indices' must have same ndim\\. Got (.+?) and (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":1049,"sourceCode":"    [[10 99 20]\n     [99 40 50]]\n  \"\"\"\n  if inplace:\n    raise ValueError(\n      \"jax.numpy.put_along_axis cannot modify arrays in-place, because JAX arrays\"\n      \"are immutable. Pass inplace=False to instead return an updated array.\")\n\n  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","sourceCodeStart":1031,"sourceCodeEnd":1067,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L1031-L1067","documentation":"put_along_axis scatters values at positions given by indices and requires arr and indices to have matching ranks so each index dimension maps to an array dimension. It fires when arr.ndim differs from indices.ndim, making the scatter ambiguous.","triggerScenarios":"jnp.put_along_axis(arr, indices, values, axis=k) where arr.ndim != indices.ndim, e.g. arr (B, N) with indices (B,) or (B, N, 1).","commonSituations":"Forgetting to append/remove a trailing axis on indices; raveling one argument but not the other before the call.","solutions":["Match ranks: indices = indices[..., None] or indices.reshape(arr.shape-structured shape)","If you meant elementwise flat update, ravel arr too when using axis=None"],"exampleFix":"// before\na = jnp.put_along_axis(a, idx, v, axis=1, inplace=False)  # idx: (B,)\n// after\na = jnp.put_along_axis(a, idx[:, None], v, axis=1, inplace=False)","handlingStrategy":"validation","validationCode":"assert arr.ndim == indices.ndim, f'ndim mismatch: arr={arr.ndim}, idx={indices.ndim}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Add trailing axes to indices: idx[..., None]","Keep arr and indices ranks aligned through pipelines"],"tags":["jax","put-along-axis","ndim","shape-mismatch"],"backgroundTag":"index-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}