JAX Skill

JAX is Autograd and XLA, brought together for high-performance machine learning research.

Contents

• [Concepts & Theory](reference.md)

- Immutability - The 4 Transformations - Pytrees

• [Code Examples](examples.md)

- jit, grad, vmap, random usage - Control Flow (scan, cond, fori_loop) - Parallelism (sharding)

Common Workflows

1. Developing a new Model

1. Define your parameters as a Pytree (dict/dataclass).
2. Define your forward pass function (pure).
3. Define your loss function.
4. Use jax.valueandgrad to get gradients.
5. Use jax.jit to speed up the update step.
6. See [examples.md](examples.md) for snippets.

2. Debugging Shapes/NaNs

1. Disable JIT: jax.config.update("jaxdisablejit", True) to debug with standard python tools.
2. Use jax.debug.print inside JITted functions.