jax.lax.pmax

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

计算在 pmapped 轴 axis_name 上 x 的 all-reduce 最大值。

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

参数:

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

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

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

返回:

与 x 形状相同的数组，表示沿轴 axis_name 执行 all-reduce max 的结果。