jax-ml/jax · error · NotImplementedError

Because JAX arrays are immutable, jnp.ufunc.at() cannot oper

Error message

Because JAX arrays are immutable, jnp.ufunc.at() cannot operate inplace like
np.ufunc.at(). Instead, you can pass inplace=False and capture the result; e.g.
>>> arr = jnp.add.at(arr, ind, val, inplace=False)

What it means

Unlike numpy, JAX arrays are immutable, so ufunc.at(..., inplace=True) cannot mutate the input in place. JAX raises NotImplementedError with guidance to use inplace=False and capture the returned updated array.

Source

Thrown at jax/_src/numpy/ufunc_api.py:442

    Examples:

      Add numbers to specified indices:

      >>> x = jnp.ones(10, dtype=int)
      >>> indices = jnp.array([2, 5, 7])
      >>> values = jnp.array([10, 20, 30])
      >>> jnp.add.at(x, indices, values, inplace=False)
      Array([ 1,  1, 11,  1,  1, 21,  1, 31,  1,  1], dtype=int32)

      This is roughly equivalent to JAX's :meth:`jax.numpy.ndarray.at` method
      called this way:

      >>> x.at[indices].add(values)
      Array([ 1,  1, 11,  1,  1, 21,  1, 31,  1,  1], dtype=int32)
    """
    if inplace:
      raise NotImplementedError(_AT_INPLACE_WARNING)

    at = self.__static_props['at'] or self._at_via_scan
    return at(a, indices) if b is None else at(a, indices, b)

  def _at_via_scan(self, a: ArrayLike, indices: Any, *args: Any) -> Array:
    assert len(args) in {0, 1}
    check_arraylike(f"{self.__name__}.at", a, *args)
    dtype = api.eval_shape(self._func, lax._one(a), *(lax._one(arg) for arg in args)).dtype
    a = lax.asarray(a).astype(dtype)
    args = tuple(lax.asarray(arg).astype(dtype) for arg in args)
    indices = indexing.eliminate_deprecated_list_indexing(indices)
    if not indices:
      return a

    shapes = [np.shape(i) for i in indices if not isinstance(i, slice)]
    shape = shapes and lax.broadcast_shapes(*shapes)
    if not shape:
      return a.at[indices].set(self(a.at[indices].get(), *args))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use inplace=False (or omit it) and reassign: arr = jnp.add.at(arr, ind, val, inplace=False)
2. Alternatively use arr.at[ind].add(val) which is the idiomatic JAX scatter

Example fix

// before
jnp.add.at(arr, ind, val, inplace=True)
// after
arr = jnp.add.at(arr, ind, val, inplace=False)
# or: arr = arr.at[ind].add(val)

updated = ufunc.at(a, indices, b, inplace=False)  # never pass inplace=True

When it happens

Trigger: jnp.add.at(arr, indices, values, inplace=True).

Common situations: Porting np.add.at(arr, idx, vals) in-place scatter-add patterns from numpy to JAX.

Related errors

• out argument of {self}
• out argument of {self.__name__}.reduce()
• out argument of {self.__name__}.accumulate()
• Value of type {type(self)} is not indexable.
• The 'out' argument to jnp.{name} is not supported.

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