Examples

NNX documentation contains two kinds of examples: example notebooks, which gradually introduce new concepts for applying NNX to different application areas, and example projects, which are more realistic representations of how nontrivial models should be implemented.

Example Notebooks

Example notebooks guide you through applying Flax models to a variety of different domains.

• Example: Using pretrained Gemma for inference with Flax NNX
• Example: Train a diffusion model for image generation
• Train a miniGPT language model with JAX

Example Projects

Example projects are hosted on the GitHub Flax repository in the examples directory.

Each example is designed to be self-contained and easily forkable, while reproducing relevant results in different areas of machine learning.

Transformers

• Gemma : A family of open-weights Large Language Model (LLM) by Google DeepMind, based on Gemini research and technology. Gemma models training and evaluation script on the One Billion Word Benchmark (LM1B).

Toy examples

NNX toy examples directory contains a few smaller, standalone toy examples for simple training scenarios.