PyPI Stats

Flax Python Packages

Python packages with the GitHub topic flax. Sorted by relevance, with stars and monthly downloads.
google
orbax-checkpoint

Orbax provides common checkpointing and persistence utilities for JAX users

11.4M 506 92
google
orbax-export

Orbax provides common checkpointing and persistence utilities for JAX users

904K 506 92
google
clu

CLU lets you write beautiful training loops in JAX.

888K 368 33
Toni-SM
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

74K 1K 141
erfanzar
eformer

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

24K 32 7
matthias-wright
flaxmodels

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

16K 264 28
erfanzar
fjformer

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

4K 32 7
Xrenya
jaxnn

Foundation models in JAX/Flax

3K 0 0
avitai
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.

3K 1 0
texttron
tevatron

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

2K 735 129
srigas
jaxkan

Kolmogorov–Arnold Networks built on JAX

2K 102 12
zlynx-ai
zlynx

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

1K 0 0
kousuke-nakano
jqmc

No description available

1K 15 1
epignatelli
helx

Interoperating between (Deep) Reiforcement Learning libraries

821 9 6
jaco-bro
diajax

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

808 29 4
evanatyourservice
image-classification-jax

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

667 3 0
darshandeshpande
jax-models

Unofficial JAX implementations of deep learning research papers

638 162 10
perrin-isir
xpag

a modular reinforcement learning library with JAX agents

498 27 6
NITHISHM2410
flax-pilot

A Simplistic trainer for Flax

496 6 0
vballoli
vit-flax

Implementation of Vision Transformers in Flax

420 18 2
tillahoffmann
flax-metrics

Flax NXX implementation of common metrics.

406 1 0
n2cholas
jax-resnet

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

374 120 8
shyamsn97
hyper-nn

Easy Hypernetworks in Pytorch and Jax

365 106 6
mlaux1
rl-blox

Modular JAX-based toolbox for implementing RL algorithms.

345 21 8