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

• Profiling JAX programs
• Device Memory Profiling
• Runtime value debugging in JAX
• GPU performance tips
• Persistent Compilation Cache

Development

• Understanding Jaxprs
• External Callbacks in JAX
• Type promotion semantics
• Pytrees

Run Time

• Ahead-of-time lowering and compilation
• JAX Errors
• Transfer guard

Custom Operations

• Pallas: a JAX kernel language