JAX ML Frameworks

Libraries and frameworks built on JAX for neural networks, optimization, and machine learning research. Includes JAX-based neural network layers, training utilities, optimization algorithms, and domain-specific extensions (graphs, replay buffers, Earth observation). Does NOT include general Python ML frameworks, probabilistic programming languages, or domain applications built with JAX.

There are 126 jax ml frameworks tracked. 5 score above 70 (verified tier). The highest-rated is explosion/thinc at 80/100 with 2,893 stars. 10 of the top 10 are actively maintained.

Get all 126 projects as JSON

curl "https://pt-edge.onrender.com/api/v1/datasets/quality?domain=ml-frameworks&subcategory=jax-ml-frameworks&limit=20"

Open to everyone — 100 requests/day, no key needed. Get a free key for 1,000/day.

#	Framework	Score	Tier	Stars	Language
1	explosion/thinc 🔮 A refreshing functional take on deep learning, compatible with your...	80	Verified	2,893	Python
2	google-deepmind/optax Optax is a gradient processing and optimization library for JAX.	79	Verified	2,207	Python
3	patrick-kidger/diffrax Numerical differential equation solvers in JAX. Autodifferentiable and...	73	Verified	1,930	Python
4	google/grain Library for reading and processing ML training data.	73	Verified	691	Python
5	patrick-kidger/equinox Elegant easy-to-use neural networks + scientific computing in JAX....	71	Verified	2,814	Python
6	extropic-ai/thrml Thermodynamic Hypergraphical Model Library in JAX	68	Established	1,026	Python
7	thomaspinder/GPJax Gaussian processes in JAX and Flax.	67	Established	601	Python
8	tumaer/JAXFLUIDS Differentiable Fluid Dynamics Package	66	Established	538	Python
9	patrick-kidger/optimistix Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox....	65	Established	553	Python
10	google-deepmind/dm-haiku JAX-based neural network library	62	Established	3,199	Python
11	MichaelTMatthews/Craftax (Crafter + NetHack) in JAX. ICML 2024 Spotlight.	62	Established	378	Python
12	google/torchax torchax is a PyTorch frontend for JAX. It gives JAX the ability to author...	61	Established	196	Python
13	google-research/kauldron Modular, scalable library to train ML models	60	Established	220	Python
14	apple/axlearn An Extensible Deep Learning Library	60	Established	2,331	Python
15	google-deepmind/kfac-jax Second Order Optimization and Curvature Estimation with K-FAC in JAX.	60	Established	317	Python
16	google/jaxopt Hardware accelerated, batchable and differentiable optimizers in JAX.	60	Established	1,030	Python
17	lockwo/distreqx Distrax, but in equinox. Lightweight JAX library of probability...	59	Established	40	Python
18	apax-hub/apax A flexible and performant framework for training machine learning potentials.	58	Established	36	Python
19	Ceyron/exponax Efficient Differentiable n-d PDE Solvers in JAX.	58	Established	163	Jupyter Notebook
20	flaport/sax S + Autograd + XLA :: S-parameter based frequency domain circuit simulations...	58	Established	110	Python
21	ekzhang/jax-js JAX in JavaScript – ML library for the web, running on WebGPU & Wasm	57	Established	733	TypeScript
22	jax-ml/bonsai Minimal, lightweight JAX implementations of popular models.	56	Established	207	Jupyter Notebook
23	n2cholas/awesome-jax JAX - A curated list of resources https://github.com/google/jax	55	Established	2,067	—
24	arpastrana/jax_fdm Auto-differentiable and hardware-accelerated force density method	53	Established	98	Python
25	e3nn/e3nn-jax jax library for E3 Equivariant Neural Networks	52	Established	224	Python
26	camail-official/discretax Discretax is a light weight collection of state space models implemented in JAX ⚡️	52	Established	32	Python
27	dpiras/cosmopower-jax Differentiable cosmological emulators: the JAX version of CosmoPower	50	Established	45	Jupyter Notebook
28	instadeepai/flashbax ⚡ Flashbax: Accelerated Replay Buffers in JAX	50	Established	274	Python
29	sotetsuk/pgx ♟️ Vectorized RL game environments in JAX	50	Established	593	Python
30	GalacticDynamics/diffraxtra Extras for Diffrax: OOP and vectorization	49	Emerging	6	Python
31	Dicklesworthstone/model_guided_research Systematic investigation of 11 exotic math frameworks (Lie groups, tropical...	49	Emerging	101	Python
32	bsc-quantic/tn4ml Tensor Networks for Machine Learning	49	Emerging	23	Python
33	thorben-frank/mlff Build neural networks for machine learning force fields with JAX	49	Emerging	133	Python
34	ergodicio/adept Automatic-Differentiation-Enabled Plasma Transport in JAX	49	Emerging	35	Python
35	gordicaleksa/get-started-with-JAX The purpose of this repo is to make it easy to get started with JAX, Flax,...	48	Emerging	776	Jupyter Notebook
36	perrin-isir/xpag a modular reinforcement learning library with JAX agents	48	Emerging	27	Python
37	poets-ai/elegy A High Level API for Deep Learning in JAX	48	Emerging	476	Python
38	danijar/ninjax General Modules for JAX	48	Emerging	72	Python
39	srush/annotated-s4 Implementation of https://srush.github.io/annotated-s4	48	Emerging	513	Python
40	google-deepmind/jeo Jeo: Jax model training lib for Earth Observation	47	Emerging	158	Python
41	genjax-community/genjax Probabilistic programming with programmable inference for parallel accelerators.	47	Emerging	47	Python
42	google/trax Trax — Deep Learning with Clear Code and Speed	46	Emerging	8,300	Python
43	MahmudulAlam/Holographic-Reduced-Representations Holographic Reduced Representations	46	Emerging	30	Python
44	BirkhoffG/jax-dataloader Pytorch-like dataloaders for JAX.	45	Emerging	102	Jupyter Notebook
45	FLAIROx/Kinetix Reinforcement learning on general 2D physics environments in JAX. ICLR 2025 Oral.	45	Emerging	234	Python
46	thebuckleylab/jpc Flexible Inference for Predictive Coding Networks in JAX.	45	Emerging	76	Python
47	eserie/wax-ml A Python library for machine-learning and feedback loops on streaming data	44	Emerging	64	Python
48	tfm000/copulax JAX based probability modelling	44	Emerging	7	Python
49	tinker495/JAxtar JAxtar is a project with a JAX-native implementation of parallelizeable A* &...	44	Emerging	44	Python
50	RobinKa/jaxga Geometric Algebra package for JAX	44	Emerging	55	Python
51	google-deepmind/jraph A Graph Neural Network Library in Jax	44	Emerging	1,467	Python
52	ikostrikov/jaxrl JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with...	44	Emerging	753	Jupyter Notebook
53	liblaf/apple 🍎 A JAX and Warp library for differentiable physics simulation, featuring...	43	Emerging	1	Python
54	shyamsn97/hyper-nn Easy Hypernetworks in Pytorch and Jax	43	Emerging	106	Jupyter Notebook
55	matomatical/hijax An introduction to vanilla JAX for deep learning research	43	Emerging	13	Python
56	matthias-wright/flaxmodels Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.	42	Emerging	265	Python
57	francois-rozet/inox Stainless neural networks in JAX	42	Emerging	34	Python
58	lockwo/awesome-jax Curated list of JAX Resources and Packages	41	Emerging	35	—
59	Ceyron/pdequinox Neural Emulator Architectures in JAX.	41	Emerging	24	Python
60	duyongan/sunstreaker 以jax为后端的类似keras的框架	40	Emerging	98	Python
61	AakashKumarNain/TF_JAX_tutorials All about the fundamental blocks of TF and JAX!	40	Emerging	279	Jupyter Notebook
62	instadeepai/catx 🐈‍⬛ Contextual bandits library for continuous action trees with smoothing in JAX	40	Emerging	71	Python
63	AaltoML/kalman-jax Approximate inference for Markov Gaussian processes using iterated Kalman...	39	Emerging	103	Jupyter Notebook
64	mancusolab/susiepca Scalable Ultra-Sparse Bayesian PCA	39	Emerging	33	Python
65	m-wojnar/reinforced-lib Reinforcement learning library	38	Emerging	65	Python
66	affjljoo3581/deit3-jax Jax/Flax implementation of DeiT and DeiT-III (ViT)	38	Emerging	19	Shell
67	cgarciae/treex A Pytree Module system for Deep Learning in JAX	38	Emerging	212	Python
68	ivy-llc/mech Mechanics functions with end-to-end support for deep learning developers,...	38	Emerging	14	Python
69	danielkelshaw/riemax Riemannian geometry in JAX	37	Emerging	16	Python
70	BobMcDear/flaim Flax Image Models - State-of-the-art pre-trained vision backbones for Flax.	37	Emerging	23	Python
71	evgenii-nikishin/omd JAX code for the paper "Control-Oriented Model-Based Reinforcement Learning...	37	Emerging	44	Jupyter Notebook
72	auto-differentiation/xad-py High-Performance Automatic Differentiation for Python	36	Emerging	19	Python
73	jeertmans/sampling-paths Generative Path Candidate Sampling for Faster Point-to-Point Ray Tracing	36	Emerging	3	Jupyter Notebook
74	IvanIsCoding/GNN-for-Combinatorial-Optimization JAX + Flax implementation of "Combinatorial Optimization with...	36	Emerging	66	Jupyter Notebook
75	XanaduAI/GradDFT GradDFT is a JAX-based library enabling the differentiable design and...	36	Emerging	112	Python
76	ericjang/pt-jax Path Tracing in JAX	35	Emerging	73	Jupyter Notebook
77	Twistient/HoloVec Holographic vectors you can compute with. Bind structure, bundle sets,...	35	Emerging	5	Python
78	satojkovic/vit-jax-flax Vision Transformer from scratch (JAX/Flax).	35	Emerging	7	Jupyter Notebook
79	wladekpal/golden-standard Is Temporal Difference Learning the Gold Standard for Stitching in RL? Code...	35	Emerging	10	Jupyter Notebook
80	google-deepmind/dks Multi-framework implementation of Deep Kernel Shaping and Tailored...	35	Emerging	76	Python
81	mila-iqia/torch_jax_interop Simple tools to mix and match PyTorch and Jax - Get the best of both worlds!	35	Emerging	37	Python
82	evgenii-nikishin/rl_with_resets JAX implementation of deep RL agents with resets from the paper "The Primacy...	34	Emerging	105	Python
83	omron-sinicx/jaxmapp JAX-based implementation for multi-agent path planning (MAPP) in continuous spaces.	34	Emerging	54	Jupyter Notebook
84	HomebrewML/revlib Simple and efficient RevNet-Library for PyTorch with XLA and DeepSpeed...	34	Emerging	132	Python
85	ml-for-gp/jaxgptoolbox Geometry processing utilities compatible with jax for autodifferentiation.	34	Emerging	90	Python
86	google-research/jestimator Amos optimizer with JEstimator lib.	33	Emerging	82	Python
87	davisyoshida/haiku-mup A port of muP to JAX/Haiku	33	Emerging	25	Python
88	phlippe/jax_trainer Lightning-like training API for JAX with Flax	33	Emerging	45	Python
89	juliuskunze/cwvae-jax Clockwork VAEs in JAX/Flax	32	Emerging	32	Python
90	amoudgl/celo Code for Celo: Training Versatile Learned Optimizers on a Compute Diet	32	Emerging	4	Python
91	yonesuke/jaxfss JAX/Flax implementation of finite-size scaling	31	Emerging	15	Python
92	cor3bit/awesome-soms A curated list of resources for second-order stochastic optimization	31	Emerging	4	—
93	graphcore-research/jax-experimental JAX for Graphcore IPU (experimental)	30	Emerging	22	Python
94	malbertosm/frp_rl Explore the "frp_rl" repository to discover the Free Random Projection...	30	Emerging	2	Python
95	ethanluoyc/magi Reinforcement learning library in JAX.	30	Emerging	101	Python
96	stefanosele/GPfY Gaussian processes with spherical harmonic features in JAX	29	Experimental	15	Python
97	pythoncrazy/jimm JAX Image Modeling of Models contains Computer Vision/Vision Language Model...	29	Experimental	7	Python
98	NITHISHM2410/flax-pilot A Simplistic trainer for Flax	29	Experimental	6	Python
99	mzguntalan/neptune [WIP] Neptune: JAX iterop-able library in Haskell.	29	Experimental	9	Haskell
100	cgarciae/nnx Neural Networks for JAX	29	Experimental	84	Python
101	norabelrose/classroom Preference-based reinforcement learning in PyTorch and JAX with a browser-based GUI.	28	Experimental	11	Python
102	mzguntalan/zephyr Zephyr is a declarative neural network library on top of JAX allowing for...	28	Experimental	36	Python
103	lweitkamp/GANs-JAX Implementation of several Generative Adversarial Networks in JAX / Flax	27	Experimental	35	Jupyter Notebook
104	alexOarga/haiku-geometric A collection of graph neural networks implementations in JAX	27	Experimental	35	Python
105	cifkao/jax-spectral Short-time Fourier transform (STFT) for JAX	27	Experimental	15	Python
106	yardenas/jax-dreamer Dreamer on JAX	27	Experimental	16	Python
107	Auxeno/ion A minimal neural network library for JAX	27	Experimental	3	Python
108	phydra-labs/phydrax Modular Physics-ML Components in JAX	26	Experimental	11	Python
109	camml-lab/reax REAX — Scalable, flexible training for JAX, inspired by the simplicity of...	25	Experimental	—	Python
110	Anuoluwapo65/pytorch-jax-implementation pytorch jax	25	Experimental	5	Jupyter Notebook
111	OleksiiBevza/jaxpsmc JAX based Preconditioned Sequential Monte Carlo framework	25	Experimental	1	Python
112	ethanluoyc/corax Corax: Core RL in JAX	23	Experimental	41	Python
113	ASEM000/serket The ✨Magical✨ JAX ML Library.	22	Experimental	18	Python
114	nathanwispinski/meta-rl A short conceptual replication of "Prefrontal cortex as a meta-reinforcement...	22	Experimental	18	Jupyter Notebook
115	forynski/jax-pid-nn High-performance JAX/Flax neural network for particle identification in...	22	Experimental	—	Jupyter Notebook
116	ysngshn/ivon-optax An Optax-based JAX implementation of the IVON optimizer for large-scale VI...	21	Experimental	9	Python
117	MizuhoAOKI/jax_playground A collection of hands-on examples for exploring numerical algorithms with JAX.	21	Experimental	5	—
118	ScottAlexanderCameron/Jynx A neural network library written in jax	21	Experimental	13	Python
119	Ceyron/trainax Training methodologies for autoregressive neural operators/emulators in JAX.	21	Experimental	10	Python
120	dtunai/LongConv-Jax Jax/Flax/Linen implementation of "Simple Hardware-Efficient Long...	19	Experimental	3	Python
121	wolfwdavid/jax-pinn JAX/Flax physics-informed neural network with jax2tf export — benchmark JAX...	14	Experimental	—	Python
122	elliotvilhelm/jax-policy-gradient 🤖 JAX implementation of REINFORCE policy gradient with baseline subtraction,...	12	Experimental	4	Python
123	PundarikakshNTripathi/JAX-Deep-Dive Deep dive into JAX: JIT compilation, autodiff, vmap, GPU neural networks,...	12	Experimental	1	Jupyter Notebook
124	bhadreshpsavani/LearningJax Learning Notes and resources on Jax	12	Experimental	7	Jupyter Notebook
125	SalamanderXing/graph-transformer-jax graph transformer in JAX	11	Experimental	—	Python
126	imdebamrita/JAX-FLAX A learning repository for JAX and FLAX, exploring automatic differentiation,...	10	Experimental	2	Jupyter Notebook

Comparisons in this category

equinox and optimistix (71 vs 65) optax and jaxopt (79 vs 60) diffrax and JAXFLUIDS (73 vs 66) diffrax and exponax (73 vs 58)