jax.experimental.multihost_utils.process_allgather#

jax.experimental.multihost_utils.process_allgather(in_tree, tiled=False)[源码]#

从各个进程收集数据。

参数:

  • in_tree (Any) – 数组的 Pytree - 每个数组在所有主机上_必须_具有相同的形状。

  • tiled (bool) – 是否堆叠或拼接输出。默认为 False，即在索引 0 处堆叠到一个新的位置轴。

返回:

Numpy 数组的 Pytrees。

  • 如果输入是一个非完全可寻址的 jax.Array，那么数据将被完全复制。

  • 如果输入是 Numpy 数组或完全可寻址的 jax.Array，则输出形状取决于 tiled 参数。如果为 False，则输出将被堆叠，否则将被拼接。

  • 如果输入是标量，则输出将被堆叠。

返回类型:

任意类型