jax-ml/jax · error · ValueError

put_along_axis argument 'values' must be broadcastable to…

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.

Solutions

Reshape/broadcast values to indices.shape: values = jnp.broadcast_to(values, indices.shape)
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)

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Pre-broadcast values to indices.shape before calling
Remember values follows the index grid, not arr's shape

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

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

Appendix: 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)