JAX Foundations

Your learning sequence

1. 1 Not solved yet.○ Pure Function vs Impure Practice · Easy
2. 2 Not solved yet.○ Tracing: Shape vs Value Practice · Medium
3. 3 Not solved yet.○ Jit with static_argnames Practice · Medium
4. 4 Not solved yet.○ Tracer Errors: Boolean Conversion and Leaks Practice · Medium
5. 5 Not solved yet.○ Functional Update with .at[].set() Practice · Easy
6. 6 Not solved yet.○ Functional Update with .at[].add() Practice · Easy
7. 7 Not solved yet.○ 2-D Scatter via .at[].add() Practice · Medium
8. 8 Not solved yet.○ JIT Compile a Function Practice · Medium
9. 9 Not solved yet.○ Vectorize with vmap Practice · Medium
10. 10 Not solved yet.○ Gradient with jax.grad Practice · Easy
11. 11 Not solved yet.○ jax.value_and_grad Practice · Easy
12. 12 Not solved yet.○ jit + grad Composition Practice · Medium
13. 13 Not solved yet.○ PRNGKey and Split Practice · Easy
14. 14 Not solved yet.○ PRNGKey fold_in Practice · Medium
15. 15 Not solved yet.○ Deterministic Batch via vmap+split Practice · Medium
16. 16 Not solved yet.○ Pytree Leaves Practice · Easy
17. 17 Not solved yet.○ Pytree Map Practice · Medium
18. 18 Not solved yet.○ Gradient of Packed Parameters Practice · Medium
19. 19 Not solved yet.○ Translate a PyTorch Matrix Pipeline to JAX Practice · Easy
20. 20 Not solved yet.○ Dtype Promotion Rules Practice · Hard
21. 21 Not solved yet.○ take_along_axis (Gather) Practice · Medium
22. 22 Not solved yet.○ Multi-Condition where Practice · Medium
23. 23 Not solved yet.○ NaN-Safe Mean (with mask) Practice · Easy
24. 24 Not solved yet.○ Explicit Broadcasting via broadcast_to Practice · Easy
25. 25 Not solved yet.○ Clip and Extrema Practice · Easy

@jax.jit
def f(x):
    print("tracing, x =", x)
    return x * 2

f(jnp.array(1.0))
f(jnp.array(3.0))
f(jnp.ones((2,)))