jax.lax.pmean

内容

jax.lax.pmean#

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

在 pmapped 轴 axis_name 上对 x 计算所有约简平均值。

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

参数：:

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

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

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

返回值：:

与 x 形状相同的数组(s)，表示沿轴 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.       ]
内容