Getting Started
jax.debug.print
jax.debug.breakpoint
checkify
Reference Documentation
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
Mesh(devices, axis_names)
Mesh
Declare the hardware resources available in the scope of this manager.
xmap(fun, in_axes, out_axes, *[, ...])
xmap
Assign a positional signature to a program that uses named array axes.