jax.experimental.pallas.atomic_cas

jax.experimental.pallas.atomic_cas(ref, cmp, val)

在引用中执行原子的比较和交换操作，将其值替换为给定的值。

参数:

• ref – 要操作的引用。

• cmp – 预期的比较值。

• val – 要交换的值。

返回:

在原子操作之前的给定索引处的值。