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#
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 tf.data and TensorFlow Datasets for scalable and reproducible input pipelines. This is the original codebase of ViT, MLP-Mixer, LiT, UViM, and many more models.