Skip to main content
Back to top
Ctrl
+
K
Flax basics
MNIST tutorial
Why Flax NNX?
Guides
Transformations
Scale up on multiple devices
Using Filters
Randomness
Performance Considerations
Evolution from Flax Linen to NNX
Use Flax NNX and Linen together
Model surgery
Save and load checkpoints
Flax NNX vs JAX Transformations
Migrating from Haiku to Flax
Example: Using Pretrained Gemma
Examples
Core examples
Flax NNX glossary
The Flax philosophy
How to contribute
API Reference
flax.nnx
graph
module
nn
Activation functions
Attention
Dtypes
Initializers
Linear
LoRA
Normalization
Recurrent
Stochastic
rnglib
spmd
state
training
Metrics
Optimizer
transforms
variables
helpers
visualization
filterlib
bridge
flax.core.frozen_dict package
flax.struct package
flax.training package
.rst
.pdf
Examples
Examples
#
Core examples
Transformers
Toy examples
