jax.distributed 模組

initialize([coordinator_address, ...])	初始化 JAX 分散式系統。
shutdown()	關閉分散式系統。