Getting Started
jax.debug.print
jax.debug.breakpoint
checkify
Reference Documentation
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
ravel_pytree(pytree)
ravel_pytree
Ravel (flatten) a pytree of arrays down to a 1D array.