jax.distributed module .. currentmodule:: jax.distributed .. automodule:: jax.distributed .. autosummary:: :toctree: _autosummary initialize shutdown