← All tracks

Flax NNX

Flax's modern object-oriented API. Stateful modules, eager-mode workflow, split/merge state, lifted transforms, training, sharding. Standalone path through the recommended API for new projects.

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