jax.lax.pmin

jax.lax.pmin(x, axis_name, *, axis_index_groups=None)

在 pmapped 轴 axis_name 上计算 x 的所有约简最小值。

如果 x 是一个 pytree，则结果等效于将此函数映射到树中的每个叶子。

参数:

• x – 具有映射轴（名为 axis_name）的数组。

• axis_name – 用于命名 pmapped 轴的可散列 Python 对象（有关详细信息，请参阅 jax.pmap() 文档）。

• axis_index_groups – 可选的列表，包含轴索引列表（例如，对于大小为 4 的轴，[[0, 1], [2, 3]] 将在头两个副本和最后两个副本上执行 pmin）。组必须完全覆盖所有轴索引一次，在 TPU 上所有组必须具有相同的大小。

返回:

与 x 形状相同的数组，表示沿轴 axis_name 进行所有约简最小值的结果。