jax.experimental.pallas.atomic_min

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

Atomically computes x_ref_or_view[idx] = min(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.