Skip to main content
Back to top
Ctrl
+
K
Flax basics
MNIST Tutorial
Guides
Transforms
Scale up on multiple devices
Using Filters
Randomness
Evolution from Linen to NNX
Use Flax NNX and Linen together
Model surgery
Flax NNX vs JAX Transformations
Migrating from Haiku to Flax
Examples
Core examples
Glossary
The Flax philosophy
How to contribute
API Reference
flax.nnx
graph
module
nn
Activation functions
Attention
Initializers
Linear
Normalization
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
