jax-ml/jax · error · ValueError

put_along_axis arguments 'arr' and 'indices' must have same…

Error message

put_along_axis arguments 'arr' and 'indices' must have same ndim. Got {arr.ndim=} and {indices.ndim=}.

What it means

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.

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

Example fix

// before
a = jnp.put_along_axis(a, idx, v, axis=1, inplace=False)  # idx: (B,)
// after
a = jnp.put_along_axis(a, idx[:, None], v, axis=1, inplace=False)

Defensive patterns

Strategy: validation

Validate before calling

assert arr.ndim == indices.ndim, f'ndim mismatch: arr={arr.ndim}, idx={indices.ndim}'

Prevention

Add trailing axes to indices: idx[..., None]
Keep arr and indices ranks aligned through pipelines

When it happens

Trigger: 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).

Common situations: Forgetting to append/remove a trailing axis on indices; raveling one argument but not the other before the call.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8396e04eea6c56fb. Report an issue: GitHub.

Appendix: source

Thrown at jax/_src/numpy/indexing.py:1049

    [[10 99 20]
     [99 40 50]]
  """
  if inplace:
    raise ValueError(
      "jax.numpy.put_along_axis cannot modify arrays in-place, because JAX arrays"
      "are immutable. Pass inplace=False to instead return an updated array.")

  arr, indices, values = util.ensure_arraylike("put_along_axis", arr, indices, values)

  original_axis = axis
  original_arr_shape = arr.shape

  if axis is None:
    arr = arr.ravel()
    axis = 0

  if not arr.ndim == indices.ndim:
    raise ValueError(
      "put_along_axis arguments 'arr' and 'indices' must have same ndim. Got "
      f"{arr.ndim=} and {indices.ndim=}."
    )

  try:
    values = util._broadcast_to(values, indices.shape)
  except ValueError:
    raise ValueError(
      "put_along_axis argument 'values' must be broadcastable to 'indices'. Got "
      f"{values.shape=} and {indices.shape=}."
    )

  idx = _make_along_axis_idx(arr.shape, indices, axis)
  result = arr.at[idx].set(values, mode=mode)

  if original_axis is None:
    result = result.reshape(original_arr_shape)

View on GitHub (pinned to 1e1c6a8fc0)