Skip to main content

Ctrl+K

This site covers the new Flax NNX API. [Click here for the old Flax Linen API]

Flax basics
MNIST tutorial
Why Flax NNX?
JAX/Flax Key Concepts
Basic Guides
Advanced Guides
Models (Bonsai)
Post-training (Tunix)
Hijax (experimental)
- Hijax
Migrating
Examples
- Example: Using pretrained Gemma for inference with Flax NNX
- Core examples
Flax NNX glossary
The Flax philosophy
How to contribute
API Reference
- flax.nnx
  - graph
  - object
  - module
  - nn
    
    Activation functions
    
    Attention
    
    Dtypes
    
    Initializers
    
    Linear
    
    LoRA
    
    Normalization
    
    Recurrent
    
    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

Basic Guides

Basic Guides#

Basic

Module & Pytree
Transformations
Model Views
Filters
Randomness
Implicit Random State
Save and load checkpoints
Flax NNX vs JAX transformations

previous

JAX/Flax Key Concepts

next

Module & Pytree

By The Flax authors

© Copyright 2023, The Flax authors.