jax.lax.pmax

目录

jax.lax.pmax#

jax.lax.pmax(x, axis_name, *, axis_index_groups=None)[源代码][源代码]#

axis_name 的 pmapped 轴上对 x 进行 all-reduce max 计算。

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

参数:
  • x – 带有映射轴名为 axis_name 的数组。

  • axis_name – 用于命名 pmapped 轴的可哈希 Python 对象(更多详情请参阅 jax.pmap() 文档)。

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

返回:

x 形状相同的数组,表示沿轴 axis_name 进行 all-reduce max 操作的结果。