jax.lax.pmean#

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

计算在pmapped轴 axis_name 上对 x 进行的全规约平均值。

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

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

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

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

返回:

形状与 x 相同的数组,表示沿轴 axis_name 的全规约平均值的结果。

例如,当有 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.       ]