jax.experimental.multihost_utils module

.. automodule:: jax.experimental.multihost_utils

.. currentmodule:: jax.experimental.multihost_utils

Multihost Utils API Reference

.. autosummary:: :toctree: _autosummary

broadcast_one_to_all sync_global_devices process_allgather assert_equal host_local_array_to_global_array global_array_to_host_local_array