JAX Logo
JAX
XLA Logo
XLA
CPU
GPU
TPU
Tracing
Compiling and
optimizations