jax.lax.pmax#
- jax.lax.pmax(x, axis_name, *, axis_index_groups=None)[source]#
在映射的轴
axis_name上计算x的全简并最大值。如果
x是一个 pytree，则结果等效于将此函数映射到树中的每个叶节点。- 参数：
x – 带有映射的轴（名为
axis_name）的数组。axis_name – 用于命名映射轴的可散列 Python 对象（有关更多详细信息，请参阅
jax.pmap()文档）。axis_index_groups – 可选的列表，其中包含轴索引的列表（例如，对于大小为 4 的轴，[[0, 1], [2, 3]] 将在头两个副本和最后两个副本上执行 pmax）。组必须完全覆盖所有轴索引一次，在 TPU 上，所有组必须具有相同的大小。
- 返回值：
与
x形状相同的数组，表示沿着轴axis_name执行全简并最大值的结果。