jax.experimental.pallas.atomic_cas

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

对 ref 中的值执行原子比较并交换操作，并将该值与给定值交换。

参数：

• ref – 要操作的 ref。

• cmp – 要比较的预期值。

• val – 要交换的 value。

返回值：

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