Libraries and Frameworks Section 040

JAX and Flax

A different way to think about numerical code: pure functions in, transformed functions out, then compiled to run fast.

10 of 10 lessons published Three reading levels on every lesson

Start with “JAX arrays and immutability”

Lessons in order

Work top to bottom. Each lesson assumes the one above it.

  1. JAX arrays and immutability
  2. jit and tracing
  3. grad and value_and_grad
  4. vmap: write the code for one example
  5. Pytrees
  6. Random numbers and PRNG keys
  7. Control flow: scan, cond and while_loop
  8. Building models with Flax
  9. Optax
  10. A full training loop in JAX