Examples

Core examples

Each example is designed to be self-contained and easily forkable, while reproducing relevant results in different areas of machine learning. These examples adhere to a shared style and functionality outlined in #231. All examples are under the folder flax/examples/. Some of the examples below have a link [Interactive] that lets you run them directly in Colab.

Image classification

  • MNIST [Interactive] : Convolutional neural network for MNIST classification (featuring simple code).

  • ImageNet : Resnet-50 on imagenet with weight decay (featuring multi host SPMD, custom preprocessing, checkpointing, dynamic scaling, mixed precision).

Reinforcement Learning

Natural language processing

Generative models

  • Variational auto-encoder : Trained on binarized MNIST (featuring simple code, vmap).

  • PixelCNN++ : Trained on cifar10 (featuring single host SPMD, checkpointing, Polyak decay).

Community Examples

In addition to the curated list of official Flax examples, there is a growing community of people using Flax to build new types of machine learning models. We are happy to showcase any example built by the community here! If you want to submit your own example, we suggest that you start by forking one of the official Flax example, and start from there.

Using Linen

Link

Author

Task type

Reference

matthias-wright/flaxmodels

@matthias-wright

Various

Various

google/vision_transformer

@andsteing

Image classification, fine-tuning

https://arxiv.org/abs/2010.11929 and https://arxiv.org/abs/2105.01601

JAX-RL

@henry-prior

Reinforcement learning

N/A

DCGAN Colab

@bkkaggle

Image Synthesis

https://arxiv.org/abs/1511.06434

BigBird Fine-tuning

@vasudevgupta7

Question-Answering

https://arxiv.org/abs/2007.14062

Using the Deprecated flax.nn API

The following examples were created using the old pre-Linen API. You can still look at them for inspiration. Or maybe ask the authors if they would accept a pull request to the new API if you want to earn some Flax karma?

Link

Author

Task type

Reference

Gaussian Processes regression

@danieljtait

Regression

N/A

DQN

@joaogui1

Reinforcement learning

N/A

Various CIFAR SOTA Models

@PForet

Image Classification

N/A

More examples

Looking for “FOO” implemented in Flax? We use GitHub issues to keep track of which models people are most interested in seeing re-implemented in Flax. If you can’t find what you’re looking for in the list “example requested”, file an issue with this template. If the model you are looking for has already been requested by others, upvote the issue to help us see which ones are the most requested.

Looking to implement something in Flax? Consider looking at the list “example requested” and go ahead and build it!