jax.distributed module

initialize([coordinator_address, ...])	Initializes the JAX distributed system.
shutdown()	Shuts down the distributed system.