Flax

Your learning sequence

1. 1 Not solved yet.○ Module with @nn.compact Practice · Medium
2. 2 Not solved yet.○ Module with setup() (alternative to compact) Practice · Medium
3. 3 Not solved yet.○ Custom Parameter Initializer Practice · Medium
4. 4 Not solved yet.○ init() and apply() Round-Trip Practice · Medium
5. 5 Not solved yet.○ nn.Sequential Composition Practice · Medium
6. 6 Not solved yet.○ Module with Multiple Named Sub-Modules Practice · Medium
7. 7 Not solved yet.○ Module That Branches on a Config Flag Practice · Medium
8. 8 Not solved yet.○ Three Levels of Module Nesting Practice · Medium
9. 9 Not solved yet.○ Multiple PRNG Streams (params, dropout) Practice · Medium
10. 10 Not solved yet.○ Train vs Eval Branches via train Flag Practice · Medium
11. 11 Not solved yet.○ Implement Dense from Scratch Practice · Medium
12. 12 Not solved yet.○ Implement Conv1D from Scratch Practice · Medium
13. 13 Not solved yet.○ Implement Conv2D with Stride and Padding Practice · Medium
14. 14 Not solved yet.○ Implement Transposed Convolution Practice · Medium
15. 15 Not solved yet.○ Implement Depthwise-Separable Convolution Practice · Hard
16. 16 Not solved yet.○ Implement LayerNorm with γ/β Practice · Medium
17. 17 Not solved yet.○ Implement BatchNorm with Mutable batch_stats Practice · Hard
18. 18 Not solved yet.○ Implement GroupNorm Practice · Medium
19. 19 Not solved yet.○ Implement RMSNorm (Modern LLM Norm) Practice · Medium
20. 20 Not solved yet.○ Implement Dropout with RNG Threading Practice · Medium
21. 21 Not solved yet.○ Scaled Dot-Product Attention Practice · Medium
22. 22 Not solved yet.○ Multi-Head Self-Attention with Flax Practice · Medium
23. 23 Not solved yet.○ Causal Multi-Head Self-Attention Practice · Medium
24. 24 Not solved yet.○ Cross-Attention with Flax MHA Practice · Medium
25. 25 Not solved yet.○ Multi-Head Attention with KV Cache Practice · Hard
26. 26 Not solved yet.○ Grouped-Query Attention (GQA) Practice · Hard
27. 27 Not solved yet.○ Multi-Query Attention (MQA) Practice · Hard
28. 28 Not solved yet.○ Sliding-Window Self-Attention Practice · Hard
29. 29 Not solved yet.○ ALiBi: Attention with Linear Biases Practice · Hard
30. 30 Not solved yet.○ Block-Diagonal Attention Mask Practice · Medium
31. 31 Not solved yet.○ Token Embedding with Flax Practice · Medium
32. 32 Not solved yet.○ Sinusoidal Position Encoding Practice · Medium
33. 33 Not solved yet.○ Learned Position Embedding Practice · Medium
34. 34 Not solved yet.○ Rotary Position Embedding (RoPE) Practice · Hard
35. 35 Not solved yet.○ ALiBi Bias Matrix Practice · Hard
36. 36 Not solved yet.○ T5 Relative Position Bucketing Practice · Hard
37. 37 Not solved yet.○ Tied Input/Output Embedding Practice · Medium
38. 38 Not solved yet.○ ViT Patch Embedding Practice · Medium
39. 39 Not solved yet.○ Transformer Encoder Block (Pre-LN) Practice · Hard
40. 40 Not solved yet.○ Transformer Decoder Block (Pre-LN) Practice · Hard
41. 41 Not solved yet.○ Pre-LN vs Post-LN Residual Pattern Practice · Medium
42. 42 Not solved yet.○ Mini GPT: Decoder-Only Language Model Practice · Hard
43. 43 Not solved yet.○ Mini BERT: Encoder-Only Hidden States Practice · Hard
44. 44 Not solved yet.○ Mini T5: Shared Embeddings Across Two Stacks Practice · Hard
45. 45 Not solved yet.○ Vision Transformer (Mean-Pool Variant) Practice · Hard
46. 46 Not solved yet.○ Vision Transformer with [CLS] Token Practice · Hard
47. 47 Not solved yet.○ DeiT-Style Transformer with Two Output Tokens Practice · Hard
48. 48 Not solved yet.○ SwiGLU Feed-Forward Network Practice · Medium
49. 49 Not solved yet.○ ResNet Basic Block Practice · Hard
50. 50 Not solved yet.○ ResNet Bottleneck Block Practice · Hard
51. 51 Not solved yet.○ Tiny ResNet Classifier Practice · Hard
52. 52 Not solved yet.○ Tiny U-Net Practice · Hard
53. 53 Not solved yet.○ GRU Cell Step Practice · Medium
54. 54 Not solved yet.○ LSTM Cell Step Practice · Medium
55. 55 Not solved yet.○ Bidirectional RNN (Flax) Practice · Medium
56. 56 Not solved yet.○ Mixture-of-Experts FFN Practice · Hard
57. 57 Not solved yet.○ Squeeze-and-Excitation Block (Flax) Practice · Medium
58. 58 Not solved yet.○ Vision-Language Fusion Practice · Medium
59. 59 Not solved yet.○ TrainState - One Step Practice · Medium
60. 60 Not solved yet.○ train_step with value_and_grad Practice · Medium
61. 61 Not solved yet.○ eval_step - Forward + Metrics Practice · Medium
62. 62 Not solved yet.○ Label-Smoothed Cross-Entropy Practice · Medium
63. 63 Not solved yet.○ Mixed-Precision Training Step Practice · Hard
64. 64 Not solved yet.○ Train with Mutable batch_stats Practice · Hard
65. 65 Not solved yet.○ Multi-Task Two-Head Loss Practice · Medium
66. 66 Not solved yet.○ Evaluate MSE in Equal Chunks Practice · Hard
67. 67 Not solved yet.○ Warmup-Cosine LR at Step Practice · Medium
68. 68 Not solved yet.○ Gradient Accumulation Step Practice · Hard
69. 69 Not solved yet.○ EMA of Parameters Practice · Medium
70. 70 Not solved yet.○ Count Parameter Arrays in a Dense Stack Practice · Medium
71. 71 Not solved yet.○ Repeat Dense Initialization Practice · Medium
72. 72 Not solved yet.○ Load Linear Weights into Flax Dense Practice · Hard
73. 73 Not solved yet.○ Pretrain, Then Train Only the Head Practice · Hard
74. 74 Not solved yet.○ Masked L2 Regularization with Adam Practice · Hard
75. 75 Not solved yet.○ Separate Learning Rates for Kernel and Bias Practice · Hard
76. 76 Not solved yet.○ Param Freezing via Grad Zeroing Practice · Hard
77. 77 Not solved yet.○ Test-Time Augmentation Aggregation Practice · Medium
78. 78 Not solved yet.○ Parameter Counts and Ideal Shard Shares Practice · Medium
79. 79 Not solved yet.○ nn.scan over an RNN cell Practice · Hard
80. 80 Not solved yet.○ nn.scan over layers Practice · Hard
81. 81 Not solved yet.○ nn.vmap with shared params Practice · Medium
82. 82 Not solved yet.○ nn.checkpoint (gradient checkpointing) Practice · Medium
83. 83 Not solved yet.○ nn.jit (Flax-aware JIT lift) Practice · Medium
84. 84 Not solved yet.○ nn.remat with checkpoint policies Practice · Hard
85. 85 Not solved yet.○ Composed lifts: nn.scan + nn.vmap Practice · Hard
86. 86 Not solved yet.○ Batched init via jax.vmap Practice · Medium
87. 87 Not solved yet.○ jax.lax.scan inside a Flax Module Practice · Hard
88. 88 Not solved yet.○ Manual Dense Ensemble Practice · Hard
89. 89 Not solved yet.○ PartitionSpec Layout Practice · Medium
90. 90 Not solved yet.○ Sharding Constraints: Values and Layout Practice · Medium
91. 91 Not solved yet.○ nn.with_partitioning Annotation Practice · Medium
92. 92 Not solved yet.○ flax.struct.dataclass - Pytree-Friendly State Practice · Medium
93. 93 Not solved yet.○ Param Surgery - Kernel Replace Practice · Medium
94. 94 Not solved yet.○ Param Surgery - Zero Last Layer Practice · Medium
95. 95 Not solved yet.○ Param Surgery - Freeze First Dense Practice · Medium
96. 96 Not solved yet.○ Partial Init - Warm-Start From Smaller Checkpoint Practice · Medium
97. 97 Not solved yet.○ Multiple Mutable Collections Practice · Medium
98. 98 Not solved yet.○ Param Sharing - One Module, Two Call Sites Practice · Medium
99. 99 Not solved yet.○ Simulate a Sharded Dense Forward Pass Practice · Hard
100. 100 Not solved yet.○ Mini LM Capstone: A Causal Forward Pass Practice · Hard

import jax
import jax.numpy as jnp
from flax import linen as nn

class BN(nn.Module):
    @nn.compact
    def __call__(self, x, *, train):
        return nn.BatchNorm(use_running_average=not train)(x)

x = jnp.arange(12, dtype=jnp.float32).reshape(4, 3)
m = BN()
v = m.init(jax.random.key(0), x, train=True)
try:
    m.apply(v, x, train=True)
except Exception as error:
    print(type(error).__name__)

out, updates = m.apply(v, x, train=True, mutable=["batch_stats"])
print(out.shape, sorted(updates.keys()))