Flax
latest
Quickstart
Overview
Installation
Examples
Guided Tour
JAX for the Impatient
Flax Basics
Annotated MNIST
How do I ...?
Managing Parameters and State
Ensembling on multiple devices
Learning Rate Scheduling
Extracting intermediate values
Model Surgery
Upgrading my Codebase to Optax
Design Notes
Dealing with Module Arguments
Lifted Transformations
Linen Design Principles
FLIPs
Additional material
The Flax Philosophy
How to Contribute
API reference
flax.linen package
flax.optim package
flax.serialization package
flax.core.frozen_dict package
flax.struct package
flax.jax_utils package
flax.traceback_util package
flax.traverse_util package
flax.training package
flax.config package
flax.error package
(deprecated)
flax.nn (deprecated)
Core: Module abstraction
Core: Additional
Linear modules
Normalization
flax.nn.BatchNorm
flax.nn.LayerNorm
flax.nn.GroupNorm
Pooling
Activation functions
Stochastic functions
Attention primitives
RNN primitives
Flax
Docs
»
flax.nn package (deprecated)
»
flax.nn.BatchNorm
Edit on GitHub
flax.nn.BatchNorm
ΒΆ
Read the Docs
v: latest
Versions
latest
Downloads
html
On Read the Docs
Project Home
Builds
Free document hosting provided by
Read the Docs
.