Flax Python Packages

orbax-checkpoint

Orbax provides common checkpointing and persistence utilities for JAX users

clu

CLU lets you write beautiful training loops in JAX.

orbax-export

Orbax provides common checkpointing and persistence utilities for JAX users

skrl

Modular Reinforcement Learning (RL) library (implemented in PyTorch, JAX, and NVIDIA Warp) with support for Gymnasium/Gym, NVIDIA Isaac Lab, MuJoCo Playground and other environments

eformer

(EasyDel Former) is a utility library designed to simplify and enhance the development in JAX

nmn

Not the neurons we want, but the neurons we need

fjformer

(EasyDel Former) is a utility library designed to simplify and enhance the development in JAX

flaxmodels

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.

calibrax

Unified benchmarking and profiling framework for the JAX scientific ML ecosystem. Timing, GPU/energy monitoring, FLOPS counting, roofline analysis, statistical testing, regression detection, and CI integration.

tevatron

Tevatron: A toolkit for learning and running deep dense retrieval models.

jax-models

Unofficial JAX implementations of deep learning research papers

image-classification-jax

Image classification in JAX with ViT, resnet, cifar10, cifar100, imagenette, and imagenet

jaxkan

Kolmogorov–Arnold Networks built on JAX

diajax

Dia-JAX: A JAX port of Dia, the text-to-speech model for generating realistic dialogue from text with emotion and tone control

jax-resnet

Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).

tn4ml

Tensor Networks for Machine Learning

vit-flax

Implementation of Vision Transformers in Flax

helx

Interoperate among reinforcement learning libraries with jax, pytorch, gym and dm_env

jqmc

jQMC code implements two real-space ab initio quantum Monte Carlo (QMC) methods. Variatioinal Monte Carlo (VMC) and lattice regularized diffusion Monte Carlo (LRDMC) methods. jQMC achieves high-performance computations especially on GPUs.

zlynx

A lightweight, modular library built on top of JAX and Flax NNX for implementing and training deep learning models.

gensbi

Generative Models for Simulation-Based Inference in JAX

flax-metrics

Flax NNX implementation of common metrics.

flax-pilot

A Simplistic trainer for Flax

jax-blox

A functional and lightweight neural network library for JAX.