jax.experimental.pallas.atomic_xchg

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

原子地交换给定值与给定索引处的值。

参数：

• x_ref_or_view – 要操作的引用。

• idx – 要使用的索引器。

• mask (Any | None | None) – 待记录。

返回值：

更新前给定索引处的值。