All modules for which code is available
- flax.configurations
- flax.core.frozen_dict
- flax.core.meta
- flax.core.scope
- flax.cursor
- flax.errors
- flax.experimental.nnx.nnx.helpers
- flax.experimental.nnx.nnx.module
- flax.experimental.nnx.nnx.nn.attention
- flax.experimental.nnx.nnx.nn.linear
- flax.experimental.nnx.nnx.nn.normalization
- flax.experimental.nnx.nnx.nn.stochastic
- flax.experimental.nnx.nnx.rnglib
- flax.experimental.nnx.nnx.spmd
- flax.experimental.nnx.nnx.transforms
- flax.experimental.nnx.nnx.variables
- flax.experimental.nnx.nnx.visualization
- flax.jax_utils
- flax.linen.activation
- flax.linen.attention
- flax.linen.batch_apply
- flax.linen.combinators
- flax.linen.initializers
- flax.linen.linear
- flax.linen.module
- flax.linen.normalization
- flax.linen.pooling
- flax.linen.recurrent
- flax.linen.spmd
- flax.linen.stochastic
- flax.linen.summary
- flax.linen.transforms
- flax.serialization
- flax.struct
- flax.traceback_util
- flax.training.checkpoints
- flax.training.common_utils
- flax.training.early_stopping
- flax.training.lr_schedule
- flax.training.train_state
- flax.traverse_util
- jax._src.nn.functions
- jax._src.ops.special