jax.experimental.pallas.atomic_add

jax.experimental.pallas.atomic_add(x_ref_or_view, idx, val, *, mask=None)

Atomically computes x_ref_or_view[idx] += val.

Parameters:

• x_ref_or_view – The ref to operate on.

• idx – The indexer to use.

• mask (Any | None | None) – TO BE DOCUMENTED.

Returns:

The value at the given index prior to the atomic operation.