Flax documentation

Flax is a neural network library and ecosystem for JAX that is designed for flexibility. Flax is in use by a growing community of researchers and engineers at Google who happily use Flax for their daily research.

For a quick introduction and short example snippets, see our README.

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
• Convert PyTorch Models to Flax
• Upgrading my Codebase to Optax
• Upgrading my Codebase to Linen
• Processing the entire Dataset

Design Notes

• Dealing with Module Arguments
• Lifted Transformations
• Linen Design Principles
• The Module lifecycle
• Should I use setup or nn.compact?
• FLIPs

Additional material

• The Flax Philosophy
• How to Contribute

API reference

• flax.linen package
  • Module
  • Init/Apply
  • Variables
  • Compact methods
  • No wrap methods
  • Profiling
  • Inspection
  • Transformations
  • Linear modules
  • Normalization
  • Pooling
  • Activation functions
  • Combinators
  • Attention primitives
  • Stochastic
  • RNN primitives
• flax.serialization package
  • State dicts
  • Serialization with MessagePack
• flax.core.frozen_dict package
• flax.struct package
• flax.jax_utils package
  • Multi device utilities
• flax.traceback_util package
  • Traceback filtering utils
• flax.traverse_util package
  • Traversal objects
  • Dict utils
  • Model parameter traversal
• flax.training package
  • Checkpoints
  • Learning rate schedules
  • Train state
  • Early Stopping
• flax.config package
• flax.error package