Flax NNX

Your learning sequence

1. 1 Not solved yet.○ NNX Module Basics Practice · Medium
2. 2 Not solved yet.○ nnx.Rngs Container Practice · Medium
3. 3 Not solved yet.○ NNX Call vs Init Practice · Medium
4. 4 Not solved yet.○ NNX Multi-Param Module Practice · Medium
5. 5 Not solved yet.○ NNX Mutability Demo Practice · Medium
6. 6 Not solved yet.○ NNX Module With Submodule Practice · Medium
7. 7 Not solved yet.○ NNX Deep Nesting Practice · Medium
8. 8 Not solved yet.○ NNX Conditional Submodule Practice · Medium
9. 9 Not solved yet.○ NNX Train/Eval Flag Practice · Medium
10. 10 Not solved yet.○ NNX Inspecting Model State Practice · Medium
11. 11 Not solved yet.○ NNX Variable vs Param Practice · Medium
12. 12 Not solved yet.○ NNX State Split/Merge Practice · Medium
13. 13 Not solved yet.○ NNX State Filter Practice · Medium
14. 14 Not solved yet.○ NNX State Transform Practice · Medium
15. 15 Not solved yet.○ NNX Pure State Update Practice · Medium
16. 16 Not solved yet.○ NNX Running Average / EMA Step Practice · Medium
17. 17 Not solved yet.○ NNX Counters and Buffers Practice · Medium
18. 18 Not solved yet.○ NNX Plain Array Attributes Practice · Medium
19. 19 Not solved yet.○ NNX Multi State Types Practice · Medium
20. 20 Not solved yet.○ NNX GraphDef Introspection Practice · Medium
21. 21 Not solved yet.○ NNX Implement Dense Practice · Medium
22. 22 Not solved yet.○ NNX Implement Conv1D Practice · Medium
23. 23 Not solved yet.○ NNX Implement Conv2D Practice · Medium
24. 24 Not solved yet.○ NNX Implement LayerNorm Practice · Medium
25. 25 Not solved yet.○ NNX Implement RMSNorm Practice · Medium
26. 26 Not solved yet.○ NNX Implement BatchNorm Practice · Hard
27. 27 Not solved yet.○ NNX Implement GroupNorm Practice · Medium
28. 28 Not solved yet.○ NNX Implement Dropout Practice · Medium
29. 29 Not solved yet.○ NNX Implement Embed Practice · Medium
30. 30 Not solved yet.○ NNX Implement Positional Embed Practice · Medium
31. 31 Not solved yet.○ NNX Scaled Dot-Product Attention Practice · Medium
32. 32 Not solved yet.○ NNX MHA From Scratch Practice · Hard
33. 33 Not solved yet.○ NNX Causal MHA Practice · Medium
34. 34 Not solved yet.○ NNX Cross-Attention Practice · Medium
35. 35 Not solved yet.○ NNX MHA With KV Cache Practice · Hard
36. 36 Not solved yet.○ NNX Grouped-Query Attention (GQA) Practice · Hard
37. 37 Not solved yet.○ NNX Multi-Query Attention (MQA) Practice · Hard
38. 38 Not solved yet.○ NNX Sliding-Window Attention Practice · Hard
39. 39 Not solved yet.○ NNX ALiBi-Style Attention Practice · Hard
40. 40 Not solved yet.○ NNX RoPE Attention Practice · Hard
41. 41 Not solved yet.○ NNX Transformer Encoder Block Practice · Hard
42. 42 Not solved yet.○ NNX Transformer Decoder Block Practice · Hard
43. 43 Not solved yet.○ NNX Mini-GPT Practice · Hard
44. 44 Not solved yet.○ NNX Mini-BERT Practice · Hard
45. 45 Not solved yet.○ NNX Vision Transformer Practice · Hard
46. 46 Not solved yet.○ NNX ViT with CLS Token Practice · Hard
47. 47 Not solved yet.○ NNX ResNet Basic Block Practice · Hard
48. 48 Not solved yet.○ NNX Tiny ResNet Classifier Practice · Hard
49. 49 Not solved yet.○ NNX Tiny U-Net Practice · Hard
50. 50 Not solved yet.○ NNX SwiGLU FFN Practice · Medium
51. 51 Not solved yet.○ NNX Optimizer Basics Practice · Medium
52. 52 Not solved yet.○ NNX Train Step Practice · Medium
53. 53 Not solved yet.○ NNX Eval Step Practice · Medium
54. 54 Not solved yet.○ NNX MultiMetric Practice · Medium
55. 55 Not solved yet.○ NNX Train with BatchNorm Practice · Hard
56. 56 Not solved yet.○ NNX Gradient Accumulation Practice · Hard
57. 57 Not solved yet.○ NNX Mixed-Precision Step Practice · Hard
58. 58 Not solved yet.○ NNX LR Schedule Practice · Medium
59. 59 Not solved yet.○ NNX Multi-Task Loss Practice · Medium
60. 60 Not solved yet.○ NNX Label Smoothing Practice · Medium
61. 61 Not solved yet.○ NNX nnx.jit Practice · Medium
62. 62 Not solved yet.○ NNX vmap over Batch Practice · Medium
63. 63 Not solved yet.○ NNX vmap over Ensemble Practice · Hard
64. 64 Not solved yet.○ NNX Scan RNN Practice · Hard
65. 65 Not solved yet.○ NNX Scan Layers Practice · Hard
66. 66 Not solved yet.○ NNX Checkpoint / Remat Practice · Medium
67. 67 Not solved yet.○ NNX Remat Policies Practice · Hard
68. 68 Not solved yet.○ NNX Composed Transforms Practice · Hard
69. 69 Not solved yet.○ NNX Batched Init Practice · Medium
70. 70 Not solved yet.○ NNX Fori Loop in Module Practice · Hard
71. 71 Not solved yet.○ NNX Bridge: Call Linen From NNX Practice · Hard
72. 72 Not solved yet.○ NNX Bridge: Call NNX From Linen Practice · Hard
73. 73 Not solved yet.○ NNX Bridge: Shared Params Across nnx and Linen Practice · Hard
74. 74 Not solved yet.○ NNX Port: Linen Dense → NNX Linear Practice · Medium
75. 75 Not solved yet.○ NNX Port: Linen MLP → NNX MLP Practice · Medium
76. 76 Not solved yet.○ NNX Port: Linen Transformer Block → NNX Practice · Hard
77. 77 Not solved yet.○ NNX Bridge: Train Mixed Linen+NNX Model Practice · Hard
78. 78 Not solved yet.○ NNX Bridge: State Translation NNX → Linen Practice · Medium
79. 79 Not solved yet.○ NNX: Import PyTorch Linear Weights Practice · Hard
80. 80 Not solved yet.○ NNX Coexistence: Parity With Linen Practice · Medium
81. 81 Not solved yet.○ NNX PartitionSpec Layout Practice · Medium
82. 82 Not solved yet.○ NNX State Sharding Specifications Practice · Medium
83. 83 Not solved yet.○ NNX Initialization State and Mesh Configuration Practice · Medium
84. 84 Not solved yet.○ NNX Data-Parallel Step Practice · Hard
85. 85 Not solved yet.○ NNX Tensor-Parallel Linear Practice · Hard
86. 86 Not solved yet.○ NNX SGD Baseline for Sharded Training Practice · Hard
87. 87 Not solved yet.○ NNX shard_map Simulation Practice · Hard
88. 88 Not solved yet.○ NNX Parameter Counts and Ideal Shares Practice · Medium
89. 89 Not solved yet.○ Repeat NNX Linear Initialization Practice · Medium
90. 90 Not solved yet.○ NNX Distributed Gradient Reduction Simulation Practice · Hard
91. 91 Not solved yet.○ NNX Surgery Replace Practice · Medium
92. 92 Not solved yet.○ NNX Surgery Zero Layer Practice · Medium
93. 93 Not solved yet.○ NNX Surgery Freeze Practice · Hard
94. 94 Not solved yet.○ NNX Surgery Add Layer Practice · Medium
95. 95 Not solved yet.○ NNX Debug Print Practice · Medium
96. 96 Not solved yet.○ NNX Debug Callback Practice · Medium
97. 97 Not solved yet.○ NNX Eager Debug Practice · Medium
98. 98 Not solved yet.○ NNX Graphdef vs State Debug Practice · Medium
99. 99 Not solved yet.○ NNX Tied Input and Output Embeddings Practice · Medium
100. 100 Not solved yet.○ NNX Mini Language Model Capstone Practice · Hard

import jax.numpy as jnp
from flax import nnx

a = nnx.Linear(2, 2, rngs=nnx.Rngs(0))
other = nnx.Linear(2, 2, rngs=nnx.Rngs(1))
b = nnx.Linear(2, 2, rngs=nnx.Rngs(0))
print(bool(jnp.array_equal(a.kernel[...], b.kernel[...])))