jax.experimental.multihost_utils
模块#
用于跨多主机同步和通信的实用工具。
Multihost Utils API 参考#
|
将数据从源主机(默认为主机 0)广播到所有其他主机。 |
|
在所有主机/设备之间创建一个屏障。 |
|
从跨进程收集数据。 |
|
验证所有主机是否具有相同的值树。 |
将主机本地值转换为全局分片的 jax.Array。 |
|
将全局 jax.Array 转换为主机本地 jax.Array。 |