jax.experimental.pallas.原子最小值

jax.experimental.pallas.atomic_min(x_ref_or_view, idx, val, *, mask=None)

原子性地计算 x_ref_or_view[idx] = min(x_ref_or_view[idx], val)。

参数:

• x_ref_or_view – 要操作的引用。

• idx – 要使用的索引器。

• mask (Any | None) – 待文档化。

返回:

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