jax.experimental.pallas.atomic_cas

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

執行 ref 中值與給定值的原子比較和交換。

參數:

• ref – 要操作的 ref。

• cmp – 要比較的預期值。

• val – 要交換的值。

返回:

原子操作之前給定索引處的值。