Skip to content
← All tracks

Flax NNX

JAX

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.

This track is written for JAX, which isn't the mode you're browsing in.

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

Check yourself

4 questions · one attempt each

These do not count toward finishing the track. They are here to catch the things that are easy to read past.

0 / 4

NNX layers take rngs= at construction. What does that buy over a global seed? 

nnx.Linear(2, 2, rngs=nnx.Rngs(0))
nnx.Linear(2, 2, rngs=nnx.Rngs(0))
nnx.Linear(2, 2, rngs=nnx.Rngs(1))
Question 1 of 4