Getting Started
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
Jaxpr(constvars, invars, outvars, eqns)
Jaxpr
ClosedJaxpr(jaxpr, consts)
ClosedJaxpr