All modules for which code is available

• flax.configurations
• flax.core.frozen_dict
• flax.linen.pooling
• flax.nnx.bridge.variables
• flax.nnx.bridge.wrappers
• flax.nnx.extract
• flax.nnx.filterlib
• flax.nnx.graphlib
• flax.nnx.helpers
• flax.nnx.module
• flax.nnx.nn.attention
• flax.nnx.nn.dtypes
• flax.nnx.nn.linear
• flax.nnx.nn.lora
• flax.nnx.nn.normalization
• flax.nnx.nn.recurrent
• flax.nnx.nn.stochastic
• flax.nnx.pytreelib
• flax.nnx.rnglib
• flax.nnx.spmd
• flax.nnx.statelib
• flax.nnx.summary
• flax.nnx.training.ema
• flax.nnx.transforms.autodiff
• flax.nnx.transforms.compilation
• flax.nnx.transforms.iteration
• flax.nnx.transforms.transforms
• flax.nnx.variablelib
• flax.nnx.visualization
• flax.struct
• flax.training.train_state
• flax.traverse_util
• jax._src.nn.functions
• jax._src.ops.special