jax.experimental.multihost_utils.broadcast_one_to_all

jax.experimental.multihost_utils.broadcast_one_to_all(in_tree, is_source=None)

将数据从源主机（默认情况下为主机 0）广播到所有其他主机。

参数:

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

• is_source (bool | None | None) – 可选的布尔值，表示调用者是否为源。只有“源主机”会提供广播的数据。如果为 None，则使用主机 0。

返回:

一个与 in_tree 匹配的 pytree，其中叶子现在都包含来自第一个主机的数据。

返回类型:

Any