{"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 requires arr and indices to have the same number of dimensions; after axis=None flattening both are 1-D, otherwise ranks must match exactly. The message shows both ndims.","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","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}