Core examples

Core examples#

Core examples 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.

Some of the examples below have a link “Interactive🕹” that lets you run them directly in Colab.

Transformers#

  • Gemma : A family of open-weights Large Language Model (LLM) by Google DeepMind, based on Gemini research and technology.

  • LM1B : Transformer encoder trained on the One Billion Word Benchmark.

  • Diffusion Models :

    A simple example of an image diffusion model using a U-Net architecture.

Toy examples#

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

Contents