Getting Started
jax.debug.print
jax.debug.breakpoint
checkify
Reference Documentation
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
initialize([coordinator_address, ...])
initialize
Initializes the JAX distributed system.
shutdown()
shutdown
Shuts down the distributed system.