Guides

• Basic Guides
  • Module & Pytree
  • Transformations
  • Model Views
  • Filters
  • Randomness
  • Implicit Random State
  • Save and load checkpoints
  • Flax NNX vs JAX transformations
• Advanced Guides
  • Scale up on multiple devices
  • Performance considerations
  • Use Flax NNX and Linen together
  • Model surgery