jax.experimental.pallas.swap#

jax.experimental.pallas.swap(x_ref_or_view, idx, val, *, mask=None, eviction_policy=None, _function_name='swap')[來源]#

交換給定索引處的值，並傳回舊值。

請參閱 load() 以瞭解引數的含義。

傳回:

在交換之前儲存在 ref 中的值。

傳回類型:

jax.Array