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 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.

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)

Solutions

1. Match ranks: indices = indices[..., None] or indices.reshape(arr.shape-structured shape)
2. 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)

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

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

• indices and arr must have the same number of dimensions; {}
• put_along_axis argument 'values' must be broadcastable to 'i
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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