Repositories that use Flax#

The following code bases use Flax and provide training frameworks and a wealth of examples. In many cases, you can also find pre-trained weights:

πŸ€— Hugging Face#

πŸ€— Hugging Face is a very popular library for building, training, and deploying state of the art machine learning models. These models can be applied on text, images, and audio. After organizing the JAX/Flax community week, they have now over 5,000 Flax/JAX models in their repository.

πŸ₯‘ DALLE Mini#

πŸ₯‘ DALLE Mini is a Transformer-based text-to-image model implemented in JAX/Flax that follows the ideas from the original DALLE paper by OpenAI.


Scenic is a codebase/library for computer vision research and beyond. Scenic’s main focus is around attention-based models. Scenic has been successfully used to develop classification, segmentation, and detection models for multiple modalities including images, video, audio, and multimodal combinations of them.

Big Vision#

Big Vision is a codebase designed for training large-scale vision models using Cloud TPU VMs or GPU machines. It is based on Jax/Flax libraries, and uses and TensorFlow Datasets for scalable and reproducible input pipelines. This is the original codebase of ViT, MLP-Mixer, LiT, UViM, and many more models.


T5X is a modular, composable, research-friendly framework for high-performance, configurable, self-service training, evaluation, and inference of sequence models (starting with language) at many scales.