jax-ml/jax · error · ValueError

put_along_axis argument 'values' must be broadcastable to 'i

Error message

put_along_axis argument 'values' must be broadcastable to 'indices'. Got {values.shape=} and {indices.shape=}.

What it means

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.

Source

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

  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)

  return result


### Indexing

def _is_integer_index(idx: Any) -> bool:
  return isinstance(idx, (int, np.integer)) and not isinstance(idx, (bool, np.bool_))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/broadcast values to indices.shape: values = jnp.broadcast_to(values, indices.shape)
2. Or use values[..., None] to add the trailing axis matching the index grid

Example fix

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

import jax.numpy as jnp
values_b = jnp.broadcast_to(values, indices.shape)  # raises early if incompatible

When it happens

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

Common situations: Passing per-row values where a full grid is needed; reusing values shaped for arr rather than for the index grid.

Related errors

• clamp requires min.shape == operand.shape or min.shape == ()
• clamp requires max.shape == operand.shape or max.shape == ()
• put_along_axis arguments 'arr' and 'indices' must have same
• x and y cannot be broadcast together.
• Cannot broadcast shape {self.shape} to layout {o.layout}

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