User guides

User guides are deeper dives into particular topics within JAX that become relevant as your JAX project matures into larger or deployed codebases.

Debugging and performance

• How to think in JAX
• Profiling computation
• Profiling device memory
• Debugging runtime values
• GPU performance tips
• Persistent compilation cache

Interfaces

• Pytrees
• Errors
• Ahead-of-time lowering and compilation
• Exporting and serialization
• Type promotion semantics
• Transfer guard

Custom operations

• Pallas: a JAX kernel language
• Foreign function interface (FFI)

Example applications

• Training a simple neural network, with tensorflow/datasets data loading
• Training a simple neural network, with PyTorch data loading
• Autobatching for Bayesian inference