jax.experimental.pallas.triton.atomic_cas

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

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

参数:

• ref – 要操作的 ref。

• cmp – 要与之比较的预期值。

• val – 要交换进去的值。

返回:

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