JAX Vectorization & Control Flow

Your learning sequence

1. 1 Not solved yet.○ vmap with in_axes Practice · Medium
2. 2 Not solved yet.○ vmap with out_axes Practice · Medium
3. 3 Not solved yet.○ vmap with None broadcasting Practice · Easy
4. 4 Not solved yet.○ Nested vmap (vmap-of-vmap) Practice · Hard
5. 5 Not solved yet.○ Cumulative Sum via lax.scan Practice · Easy
6. 6 Not solved yet.○ Running Mean via lax.scan Practice · Medium
7. 7 Not solved yet.○ Scan over Layer Stack Practice · Medium
8. 8 Not solved yet.○ Training Loop via lax.scan Practice · Hard
9. 9 Not solved yet.○ Scan with Per-Step Outputs Practice · Medium
10. 10 Not solved yet.○ Newton's Method via lax.while_loop Practice · Medium
11. 11 Not solved yet.○ Bounded Search via lax.while_loop Practice · Hard
12. 12 Not solved yet.○ lax.while_loop vs Python while Practice · Medium
13. 13 Not solved yet.○ x^n via lax.fori_loop Practice · Easy
14. 14 Not solved yet.○ fori vs scan vs while Practice · Medium
15. 15 Not solved yet.○ Binary Classification via lax.cond Practice · Easy
16. 16 Not solved yet.○ Multi-Branch Dispatch via lax.switch Practice · Medium
17. 17 Not solved yet.○ vmap over lax.cond Practice · Hard
18. 18 Not solved yet.○ vmap over lax.scan Practice · Hard
19. 19 Not solved yet.○ jit of scan Practice · Medium
20. 20 Not solved yet.○ vmap-cond vs where: cost tradeoff Practice · Hard
21. 21 Not solved yet.○ Dynamic Slice Practice · Medium
22. 22 Not solved yet.○ Dynamic Update Slice Practice · Medium
23. 23 Not solved yet.○ Associative Scan (Parallel Cumsum) Practice · Hard
24. 24 Not solved yet.○ lax.scan with reverse=True Practice · Medium
25. 25 Not solved yet.○ lax.map vs vmap Practice · Medium

def step(c, x):
    return c + x, c

carry, ys = jax.lax.scan(step, 0.0, jnp.arange(5.0))