JAX Foundations

Functional ML — pure functions, vmap, pytrees. The JAX way of thinking.

1. 1. ○ Vectorize with vmap
2. 2. ○ Pure Function vs Impure
3. 3. ○ Tracing: Shape vs Value
4. 4. ○ JIT Compile a Function
5. 5. ○ Jit with static_argnames
6. 6. ○ Tracer Leak Detection
7. 7. ○ Functional Update with .at[].set()
8. 8. ○ Functional Update with .at[].add()
9. 9. ○ 2-D Scatter via .at[].add()
10. 10. ○ Gradient with jax.grad
11. 11. ○ jax.value_and_grad
12. 12. ○ jit + grad Composition
13. 13. ○ PRNGKey and Split
14. 14. ○ PRNGKey fold_in
15. 15. ○ Deterministic Batch via vmap+split
16. 16. ○ Pytree Leaves
17. 17. ○ Pytree Map
18. 18. ○ Gradient over Pytree Params
19. 19. ○ JAX numpy vs PyTorch ops
20. 20. ○ Dtype Promotion Rules
21. 21. ○ take_along_axis (Gather)
22. 22. ○ Multi-Condition where
23. 23. ○ NaN-Safe Mean (with mask)
24. 24. ○ Explicit Broadcasting via broadcast_to
25. 25. ○ Clip and Extrema