jax.experimental.pallas.swap
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
