jax.lax.pmean

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

在 axis_name 的 pmapped 轴上对 x 进行 all-reduce 均值计算。

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

参数:

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

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

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

返回:

与 x 形状相同的数组，表示沿着轴 axis_name 进行 all-reduce 均值运算的结果。

例如，如果有4个可用的XLA设备：

>>> x = np.arange(4)
>>> y = jax.pmap(lambda x: jax.lax.pmean(x, 'i'), axis_name='i')(x)
>>> print(y)
[1.5 1.5 1.5 1.5]
>>> y = jax.pmap(lambda x: x / jax.lax.pmean(x, 'i'), axis_name='i')(x)
>>> print(y)
[0.        0.6666667 1.3333334 2.       ]