jax-ml/jax · error · ValueError

jax.numpy.put_along_axis cannot modify arrays in-place, beca

Error message

jax.numpy.put_along_axis cannot modify arrays in-place, because JAX arraysare immutable. Pass inplace=False to instead return an updated array.

What it means

JAX arrays are immutable, so put_along_axis cannot write in-place despite exposing an inplace flag for NumPy API compatibility. With inplace=True it raises ValueError telling you to use inplace=False.

Source

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

    - :func:`jax.numpy.place`: place elements into an array via boolean mask.
    - :func:`jax.numpy.ndarray.at`: array updates using NumPy-style indexing.
    - :func:`jax.numpy.take`: extract values from an array at given indices.
    - :func:`jax.numpy.take_along_axis`: extract values from an array along an axis.

  Examples:
    >>> from jax import numpy as jnp
    >>> a = jnp.array([[10, 30, 20], [60, 40, 50]])
    >>> i = jnp.argmax(a, axis=1, keepdims=True)
    >>> print(i)
    [[1]
     [0]]
    >>> b = jnp.put_along_axis(a, i, 99, axis=1, inplace=False)
    >>> print(b)
    [[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=}."
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass inplace=False and use the returned array: arr = jnp.put_along_axis(arr, idx, vals, axis, inplace=False)
2. Remember all JAX updates are functional — always rebind the result

Example fix

// before
jnp.put_along_axis(a, i, 99, axis=1, inplace=True)
// after
a = jnp.put_along_axis(a, i, 99, axis=1, inplace=False)

assert not inplace, 'JAX is immutable; use inplace=False'

When it happens

Trigger: Calling jnp.put_along_axis(arr, indices, values, axis, inplace=True) (the NumPy default).

Common situations: Porting np.put_along_axis calls verbatim; users assuming the NumPy signature works identically in JAX.

Related errors

• put_along_axis arguments 'arr' and 'indices' must have same
• put_along_axis argument 'values' must be broadcastable to 'i
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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