JAX

1. Key concepts

   Arrays, transformations, tracing, and the functional programming model.

2. Just-in-time compilation

   jax.jit, XLA compilation, static arguments, and tracing behavior.

3. Automatic differentiation

   grad, value_and_grad, jvp, vjp, and higher-order derivatives.

4. Automatic vectorization

   vmap for batching functions without manual loops.

5. Pytrees

   Nested containers of arrays and tree utilities.

6. Pseudorandom numbers

   Explicit PRNG keys, splitting, and reproducibility.

7. Stateful computations

   Handling state in a pure functional style.

8. Common gotchas

   Pure functions, in-place updates, control flow, and dynamic shapes.

9. Distributed arrays and sharding

   Automatic and explicit parallelism across devices with Mesh and PartitionSpec.

10. Control flow

    cond, scan, fori_loop, and while_loop under jit.

11. Debugging

    jax.debug tools, NaN checks, and inspecting traced values.

12. Pallas kernels

    Writing custom kernels for GPU and TPU.

Each section links to the official page. Titles and summaries by Extra Credit; the docs belong to their authors.