Skip to main content
Ctrl+K

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 peformance tips

Development

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

Run Time

  • Ahead-of-time lowering and compilation
  • JAX Errors
  • Transfer guard
  • Pallas: a JAX kernel language

previous

Stateful Computations in JAX

next

Profiling JAX programs