PyRank
  • Insights
  • PyPI
  • GitHub
  • Search
  • Compare
  • Advisories
  • Ecosystem
  • About

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

8.4M 507 92
google
clu

CLU lets you write beautiful training loops in JAX.

914K 366 33
google
orbax-export

Orbax provides common checkpointing and persistence utilities for JAX users

848K 507 92
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

77K 1K 142
erfanzar
eformer

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

23K 32 7
matthias-wright
flaxmodels

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

13K 264 28
erfanzar
fjformer

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

5K 32 7
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.

4K 1 0
texttron
tevatron

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

1K 736 129
evanatyourservice
image-classification-jax

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

1K 3 0
zlynx-ai
zlynx

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

1K 2 0
darshandeshpande
jax-models

Unofficial JAX implementations of deep learning research papers

1K 162 10
kousuke-nakano
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.

981 16 2
epignatelli
helx

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

974 9 6
srigas
jaxkan

Kolmogorov–Arnold Networks built on JAX

905 104 12
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

740 29 4
n2cholas
jax-resnet

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

693 120 8
vballoli
vit-flax

Implementation of Vision Transformers in Flax

667 18 2
Xrenya
jaxnn

JaxNN: Foundation Models in JAX/Flax

536 0 0
NITHISHM2410
flax-pilot

A Simplistic trainer for Flax

446 6 0
tillahoffmann
flax-metrics

Flax NXX implementation of common metrics.

417 1 0
perrin-isir
xpag

a modular reinforcement learning library with JAX agents

386 27 6
DBraun
jraphx

Graph Neural Networks with JAX/NNX

383 13 1
shyamsn97
hyper-nn

Easy Hypernetworks in Pytorch and Jax

303 106 6
    • Data from PyPI, GitHub, ClickHouse, and BigQuery