Catalyst — Peer-reviewed • Rank 19.7 • Science 100%

Catalyst: a Python JIT compiler for auto-differentiable hybrid quantum programs - Published in JOSS (2024)

pyhf — Peer-reviewed • Rank 19.4 • Science 100%

pyhf: pure-Python implementation of HistFactory statistical models - Published in JOSS (2021)

Dynamax — Peer-reviewed • Rank 18.0 • Science 100%

Dynamax: A Python package for probabilistic state space modeling with JAX - Published in JOSS (2025)

Foolbox Native — Peer-reviewed • Rank 20.0 • Science 95%

Foolbox Native: Fast adversarial attacks to benchmark the robustness of machine learning models in PyTorch, TensorFlow, and JAX - Published in JOSS (2020)

mpi4jax — Peer-reviewed • Rank 14.6 • Science 100%

mpi4jax: Zero-copy MPI communication of JAX arrays - Published in JOSS (2021)

qujax — Peer-reviewed • Rank 13.9 • Science 98%

qujax: Simulating quantum circuits with JAX - Published in JOSS (2023)

Scientific Computational Imaging Code (SCICO) — Peer-reviewed • Rank 13.0 • Science 98%

Scientific Computational Imaging Code (SCICO) - Published in JOSS (2022)

flowMC — Peer-reviewed • Rank 15.8 • Science 95%

flowMC: Normalizing flow enhanced sampling package for probabilistic inference in JAX - Published in JOSS (2023)

HOI — Peer-reviewed • Rank 10.6 • Science 100%

HOI: A Python toolbox for high-performance estimation of Higher-Order Interactions from multivariate data - Published in JOSS (2024)

CR-Sparse — Peer-reviewed • Rank 9.2 • Science 98%

CR-Sparse: Hardware accelerated functional algorithms for sparse signal processing in Python using JAX - Published in JOSS (2021)

unxt — Peer-reviewed • Rank 13.2 • Science 93%

unxt: A Python package for unit-aware computing with JAX - Published in JOSS (2025)

SGMCMCJax — Peer-reviewed • Rank 11.7 • Science 93%

SGMCMCJax: a lightweight JAX library for stochastic gradient Markov chain Monte Carlo algorithms - Published in JOSS (2022)

jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs • Rank 12.0 • Science 92%

jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs - Published in JOSS (2026)

jaxKAN — Peer-reviewed • Rank 10.4 • Science 93%

jaxKAN: A unified JAX framework for Kolmogorov-Arnold Networks - Published in JOSS (2025)

SBIAX — Peer-reviewed • Rank 8.1 • Science 93%

SBIAX: Density-estimation simulation-based inference in JAX - Published in JOSS (2025)

PhotonWeave — Peer-reviewed • Rank 7.2 • Science 93%

PhotonWeave - Published in JOSS (2025)

veros • Rank 14.5 • Science 85%

The versatile ocean simulator, in pure Python, powered by JAX.

MDPax: GPU-accelerated MDP solvers in Python with JAX • Rank 6.0 • Science 93%

MDPax: GPU-accelerated MDP solvers in Python with JAX - Published in JOSS (2025)

deepmd-kit • Rank 21.7 • Science 77%

A deep learning package for many-body potential energy representation and molecular dynamics

GPJax — Peer-reviewed • Rank 17.6 • Science 77%

GPJax: A Gaussian Process Framework in JAX - Published in JOSS (2022)

tensorcircuit • Rank 15.5 • Science 77%

Tensor network based quantum software framework for the NISQ era

tensorcircuit-ng • Rank 12.6 • Science 77%

The next-gen tensor network based quantum software framework: superseding the original TensorCircuit

deepxde • Rank 12.6 • Science 77%

A library for scientific machine learning and physics-informed learning

tensorwaves • Rank 12.3 • Science 77%

Python fitter package for multiple computational back-ends

smart-transformers • Rank 7.9 • Science 77%

Smart Transformers are a versatile machine learning tool that can be integrated with Pytorch, TensorFlow, and JAX. Smart transformers provide accurate computations required for cryptographic algorithms. These transformers is that they are independent modules, making it efficient to experiment with various research projects related to cryptanalysis

kubeflow-katib • Rank 20.7 • Science 64%

Automated Machine Learning on Kubernetes

https://github.com/adaptive-intelligent-robotics/qdax • Rank 16.5 • Science 67%

Accelerated Quality-Diversity

relaxed • Rank 9.8 • Science 72%

Provides differentiable versions of common HEP operations and objectives.

felupe • Rank 13.8 • Science 67%

:mag: finite element analysis for continuum mechanics of solid bodies

ivy • Rank 26.5 • Science 54%

Convert Machine Learning Code Between Frameworks

iree-base-compiler • Rank 25.0 • Science 54%

A retargetable MLIR-based machine learning compiler and runtime toolkit.

SciMLBenchmarks • Rank 11.5 • Science 67%

Scientific machine learning (SciML) benchmarks, AI for science, and (differential) equation solvers. Covers Julia, Python (PyTorch, Jax), MATLAB, R

pycharge • Rank 11.3 • Science 67%

Differentiable electromagnetics simulation library for moving point charges built on JAX

tensorly • Rank 23.4 • Science 54%

TensorLy: Tensor Learning in Python.

zea • Rank 12.3 • Science 64%

zea: A Toolbox for Cognitive Ultrasound Imaging

mosec • Rank 20.4 • Science 54%

A high-performance ML model serving framework, offers dynamic batching and CPU/GPU pipelines to fully exploit your compute machine

rapidae • Rank 6.6 • Science 67%

Explore, compare and develop autoencoder models with a back-end agnostic framework

https://github.com/ami-iit/jaxsim • Rank 14.5 • Science 59%

A differentiable physics engine and multibody dynamics library for control and robot learning.

s2fft • Rank 14.1 • Science 59%

S2FFT: Differentiable and accelerated spherical transforms

netket • Rank 19.1 • Science 54%

Machine learning algorithms for many-body quantum systems

einops • Rank 28.9 • Science 44%

Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

pennylane • Rank 25.3 • Science 46%

PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Built by researchers, for research.

jax • Rank 17.2 • Science 54%

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

pyhpc-benchmarks • Rank 6.9 • Science 64%

A suite of benchmarks for CPU and GPU performance of the most popular high-performance libraries for Python :rocket:

vaes • Rank 1.8 • Science 67%

Reproducible code showing the various types of variational autoencoders I have implemented

https://github.com/google/trax • Rank 20.8 • Science 46%

Trax — Deep Learning with Clear Code and Speed

thinc • Rank 12.2 • Science 54%

🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

https://github.com/cbg-ethz/scatrex • Rank 5.9 • Science 57%

Map single-cell transcriptomes to copy number evolutionary trees.

gwkokab • Rank 8.8 • Science 54%

An open-source JAX-based gravitational-wave population inference toolkit for parametric models

dawdreamer • Rank 16.1 • Science 46%

Digital Audio Workstation with Python; VST instruments/effects, parameter automation, FAUST, JAX, Warp Markers, and JUCE processors

pychastic • Rank 7.2 • Science 54%

Pychastic is a stochastic differential equations integrator written entirely in python.

opennsfw2 • Rank 17.2 • Science 44%

Keras implementation of the Yahoo Open-NSFW model

turbozero • Rank 5.7 • Science 54%

fast + parallel AlphaZero in JAX

gradients • Rank 5.3 • Science 54%

Build your deep learning algorithms with confidence

https://github.com/pyro-ppl/numpyro • Rank 12.6 • Science 46%

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

ott-jax • Rank 20.1 • Science 36%

Optimal transport tools implemented with the JAX framework, to solve large scale matching problems of any flavor.

https://github.com/astro-informatics/s2wav • Rank 8.7 • Science 46%

Differentiable and accelerated wavelet transform on the sphere with JAX

scimlbenchmarksoutput • Rank 5.4 • Science 49%

SciML-Bench Benchmarks for Scientific Machine Learning (SciML), Physics-Informed Machine Learning (PIML), and Scientific AI Performance

https://github.com/google/etils • Rank 26.4 • Science 26%

Collection of eclectic utils for python.

jaxoplanet • Rank 12.0 • Science 39%

Astronomical time series analysis with JAX

pygrpy • Rank 5.7 • Science 44%

Python port of Generalized Rotne Prager Yamakawa hydrodynamic tensors.

audiotree • Rank 5.4 • Science 44%

Audio data loading and augmentations in JAX

https://github.com/ami-iit/adam • Rank 13.1 • Science 36%

adam implements a collection of algorithms for calculating rigid-body dynamics in Jax, CasADi, PyTorch, and Numpy.

pyprobml • Rank 13.0 • Science 36%

Python code for "Probabilistic Machine learning" book by Kevin Murphy

https://github.com/google-research/scenic • Rank 12.8 • Science 36%

Scenic: A Jax Library for Computer Vision Research and Beyond

jaxampler • Rank 7.2 • Science 41%

An open-source JAX-based statistical sampling toolkit 🧪

https://github.com/google-research/jestimator • Rank 10.8 • Science 36%

Amos optimizer with JEstimator lib.

dx7-jax • Rank 2.7 • Science 44%

Yamaha DX7 synthesizer with JAX

https://github.com/awslabs/fortuna • Rank 9.1 • Science 36%

A Library for Uncertainty Quantification.

https://github.com/google-deepmind/dm_pix • Rank 8.9 • Science 36%

PIX is an image processing library in JAX, for JAX.

yax • Rank 0.0 • Science 44%

Yet Another X: JAX/FLAX module tracing, modification, and evaluation.

https://github.com/cbg-ethz/scdef • Rank 7.8 • Science 36%

Deep exponential families for single-cell data.

https://github.com/google-research/sofima • Rank 17.4 • Science 26%

Scalable Optical Flow-based Image Montaging and Alignment

meent • Rank 11.9 • Science 31%

Electromagnetic simulation (RCWA) & optimization package in Python

https://github.com/google-deepmind/pgmax • Rank 6.4 • Science 36%

Loopy belief propagation for factor graphs on discrete variables in JAX

https://github.com/google-deepmind/jeo • Rank 5.9 • Science 36%

Jeo: Jax model training lib for Earth Observation

https://github.com/bminixhofer/tokenkit • Rank 3.5 • Science 36%

A toolkit implementing advanced methods to transfer models and model knowledge across tokenizers.

https://github.com/dbraun/dac-jax • Rank 3.2 • Science 36%

JAX Implementations of Descript Audio Codec and EnCodec

https://github.com/cbg-ethz/jnotype • Rank 2.5 • Science 36%

Probabilistic modeling of high-dimensional binary data in JAX

helx • Rank 8.5 • Science 26%

Interoperating between (Deep) Reiforcement Learning libraries

https://github.com/codingfisch/viiew • Rank 4.7 • Science 26%

View(your_data) 🧐 Supports numpy, pandas, torch, jax, tinygrad...

jaxwt • Rank 11.3 • Science 18%

Differentiable and gpu enabled fast wavelet transforms in JAX.

https://github.com/tummfm/jax-dimenet • Rank 5.5 • Science 23%

Jax / Haiku implementation of DimeNet++.

https://github.com/alpa-projects/alpa • Rank 18.5 • Science 10%

Training and serving large-scale neural networks with auto parallelization.

https://github.com/askabalan/jax-hpc-profiler • Rank 6.7 • Science 13%

A convinient python script to plot performance graphs for HPC JAX codes

https://github.com/eserie/wax-ml • Rank 8.9 • Science 10%

A Python library for machine-learning and feedback loops on streaming data

https://github.com/nicholas-miklaucic/eins • Rank 5.0 • Science 13%

One tensor operation is all you need

https://github.com/beomi/gemma-easylm • Rank 6.3 • Science 10%

Train GEMMA on TPU/GPU! (Codebase for training Gemma-Ko Series)

https://github.com/ivy-llc/demo-utils • Rank 0.7 • Science 13%

A set of utilities for creating visual demos for Ivy libraries.

https://github.com/4ment/phylojax • Rank 0.7 • Science 13%

https://github.com/zceemja/si-base • Rank 2.9 • Science 10%

A basic python package that converts numerical strings with units to base units

AutoPDEx — Peer-reviewed • Science 93%

AutoPDEx: An Automized Partial Differential Equation solver based on JAX - Published in JOSS (2025)

pyrddlgym-jax • Science 67%

JAX compilation of RDDL description files, and a differentiable planner in JAX.

https://github.com/instadeepai/jumanji • Science 36%

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

https://github.com/jiayaobo/fenbux • Science 13%

A Simple Statistical Distribution Library in JAX

https://github.com/jiayaobo/stamox • Science 13%

make your statistical research faster

https://github.com/ceyron/exponax • Science 36%

Efficient Differentiable n-d PDE solvers in JAX.

diff-ml • Science 44%

Second-Order Differential ML

https://github.com/apax-hub/apax • Science 49%

A flexible and performant framework for training machine learning potentials.

swarmrl • Science 26%

Multi agent reinforcement learning for intelligent active matter