jax.experimental.pallas.atomic_xchg

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

Atomically exchanges the given value with the value at the given index.

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 aupdate.