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

substrax

JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging.

datarax

Differentiable data pipelines for JAX/Flax NNX: HuggingFace, TFDS and ArrayRecord sources, augmentation stages, scan-based epochs, exact mid-epoch resume.

avitai-artifex

Generative modeling for JAX/Flax NNX: VAEs, GANs, DDPM/DDIM/score/DiT/latent diffusion, normalizing flows, energy-based, autoregressive and geometric (SE(3), protein, point cloud, mesh) models; image, text, audio, molecular, protein, tabular and time-series modalities; a shared Trainer with callbacks and extensions; evaluation through calibrax.

eformer

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

fjformer

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

dantinox

DantinoX: A modular, memory-efficient Transformer implementation in JAX/Flax NNX. Includes Sparse MoE, GQA, Sliding Window Attention, Gradient Accumulation and Checkpointing

nmn

Not the neurons we want, but the neurons we need

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.

helx

Interoperating between (Deep) Reiforcement Learning libraries

diffav

Physics-informed, RL-aligned evaluation engine for autonomous driving: JAX/Flax NNX trajectory-diffusion world models on the Waymo Open Dataset, DPO and reward-guided adversarial steering toward the scenarios that break AV stacks, physics feasibility, occupancy flow, WOSAC-style realism metrics and sensor simulation.

tevatron

Tevatron - Unified Document Retrieval Toolkit across Scale, Language, and Modality. Demo in SIGIR 2023, SIGIR 2025.

flax-pilot

A Simplistic trainer for Flax

librosax

Librosa in JAX

dewml

A nice ML framework based on Jax/Flax. Based on https://github.com/AshishKumar4/FlaxDiff

gensbi

Generative Models for Simulation-Based Inference in JAX

flaxmodels

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

jestimator

Amos optimizer with JEstimator lib.

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.

rl-blox

Modular JAX-based toolbox for implementing RL algorithms.