JAX Autodiff

Your learning sequence

1. 1 Not solved yet.○ jvp Basics Practice · Easy
2. 2 Not solved yet.○ jacfwd vs jacrev Practice · Medium
3. 3 Not solved yet.○ jvp for Sensitivity Analysis Practice · Medium
4. 4 Not solved yet.○ vjp Basics Practice · Easy
5. 5 Not solved yet.○ Build a Jacobian with VJPs Practice · Medium
6. 6 Not solved yet.○ grad vs vjp Practice · Medium
7. 7 Not solved yet.○ Hessian of a Quadratic Practice · Easy
8. 8 Not solved yet.○ HVP via grad-of-grad Practice · Medium
9. 9 Not solved yet.○ HVP via jvp-of-grad Practice · Medium
10. 10 Not solved yet.○ Custom VJP: Stable log1pexp Practice · Medium
11. 11 Not solved yet.○ Custom JVP: Clip with Pass-through Gradient Practice · Medium
12. 12 Not solved yet.○ Custom VJP: Implicit Function Theorem Practice · Hard
13. 13 Not solved yet.○ stop_gradient: Target Network Practice · Easy
14. 14 Not solved yet.○ Straight-Through Estimator Practice · Medium
15. 15 Not solved yet.○ Gradient Checkpointing: Basics Practice · Medium
16. 16 Not solved yet.○ Checkpoint with Save Policy Practice · Hard
17. 17 Not solved yet.○ Checkpointed Deep Stack via scan Practice · Medium
18. 18 Not solved yet.○ Per-Example Gradients via vmap(grad(...)) Practice · Medium
19. 19 Not solved yet.○ vmap(grad) vs grad(sum(vmap)) Practice · Medium
20. 20 Not solved yet.○ Microbatched Gradient Accumulation via scan Practice · Hard
21. 21 Not solved yet.○ jax.linearize Primitive Practice · Medium
22. 22 Not solved yet.○ Jacobian with jacrev Practice · Hard
23. 23 Not solved yet.○ Higher-Order custom_vjp Practice · Hard
24. 24 Not solved yet.○ Saved Residuals in custom_vjp Practice · Medium
25. 25 Not solved yet.○ Grad through stop_gradient Practice · Medium

jax.value_and_grad(lambda x: (x ** 2).sum())(jnp.array([3.0]))