jax — open-source software health

PyPI · Go · crates.io

94Excellenthealth index

The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.

Python · Go★ 11.2K↓ 23M/moJul 20, 2026

MITJul 20, 2026 · metrics 1.13.0

PyPI

86Excellenthealth index

keras-team/keras

Deep Learning for humans

Python★ 64.2K↓ 18.4M/moJul 16, 2026

Apache-2.0Jul 16, 2026 · metrics 1.13.0

PyPI

83Goodhealth index

jax-ml/jax

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

Python★ 36KJul 16, 2026

Apache-2.0Jul 16, 2026 · metrics 1.13.0

PyPI

82Goodhealth index

HIPS/autograd

Efficiently computes derivatives of NumPy code.

Python★ 7,512Jul 18, 2026

MITJul 18, 2026 · metrics 1.13.0

PyPI

81Goodhealth index

google/grain

Library for reading and processing ML training data.

Python★ 752Jul 21, 2026

Apache-2.0Jul 21, 2026 · metrics 1.13.0

PyPI · crates.io

78Goodhealth index

google/flax

Flax is a neural network library for JAX that is designed for flexibility.

Jupyter Notebook · Python★ 7,271↓ 4.3M/moJul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

PyPI

75Goodhealth index

tensorflow/datasets

TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...

Python★ 4,582Jul 21, 2026

Apache-2.0Jul 21, 2026 · metrics 1.13.0

PyPI

74Goodhealth index

google/etils

Collection of eclectic utils for python.

Python · Jupyter Notebook★ 257↓ 8.3M/moJul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

PyPI

72Goodhealth index

arogozhnikov/einops

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

Python · Jupyter Notebook★ 9,548Jul 17, 2026

MITJul 17, 2026 · metrics 1.13.0

PyPI · npm

72Goodhealth index

magenta/magenta-realtime

Magenta RealTime 2: An Open-Weights Live Music Model

Python · C++★ 1,683↓ 817/moJul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

PyPI

71Goodhealth index

google-deepmind/chex

No repository description published.

Python★ 948↓ 2M/moJul 21, 2026

Apache-2.0Jul 21, 2026 · metrics 1.13.0

PyPI

71Goodhealth index

google/torchax

torchax is a PyTorch frontend for JAX. It gives JAX the ability to author JAX programs using familiar PyTorch syntax. It also provides JAX-Pytorch interoperability, meaning, one can mix JAX & Pytorch syntax together when authoring ML programs, and run it in every hardware JAX can run.

Python★ 231↓ 116.4K/moJul 19, 2026

Apache-2.0Jul 19, 2026 · metrics 1.13.0

PyPI

71Goodhealth index

patrick-kidger/jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/

Python★ 1,842↓ 6.4M/moJul 18, 2026

Custom licenseJul 18, 2026 · metrics 1.13.0

PyPI

68Moderatehealth index

galacticdynamics/phasecurvefit

Construct Paths through Phase-Space Points

Python★ 4Jul 17, 2026

MITJul 17, 2026 · metrics 1.13.0

PyPI

65Moderatehealth index

tillahoffmann/jax-mps

A JAX backend for Apple Metal Performance Shaders (MPS), enabling GPU-accelerated JAX computations on Apple Silicon.

C++ · Python★ 192↓ 3,492/moJul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

PyPI

62Moderatehealth index

tenax-lab/tenax

JAX-based tensor network library with label-based contraction and symmetry-aware block-sparse tensors

Python★ 7Jul 19, 2026

Apache-2.0Jul 19, 2026 · metrics 1.13.0

PyPI

62Moderatehealth index

tensorcircuit/tensorcircuit-ng

Next-gen AI-native tensor-network-based quantum software framework

Python★ 81↓ 0/moJul 13, 2026

Apache-2.0Jul 13, 2026 · metrics 1.13.0

PyPI · npm

61Moderatehealth index

explosion/thinc

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

Python★ 2,890↓ 22.6M/moJul 18, 2026

MITJul 18, 2026 · metrics 1.13.0

PyPI

60Moderatehealth index

elkins-lab/diff-biophys

Differentiable biophysical modeling in JAX (SAXS, NMR, CD), with Jupyter Notebook examples.

Python★ 0Jul 17, 2026

MITJul 17, 2026 · metrics 1.13.0

PyPI

56Moderatehealth index

elkins-lab/diff-ensemble

Differentiable VAE framework for predicting protein structural ensembles (IDPs) consistent with SAXS and NMR data. Built on JAX/Flax.

Python★ 1↓ 158/moJul 17, 2026

MITJul 17, 2026 · metrics 1.13.0

PyPI

56Moderatehealth index

elkins-lab/resonance-flow

JAX-native differentiable protein folding framework integrating experimental NMR constraints (RDCs) and biophysical "self-correction." Several Jupyter Notebooks visualize the concepts.

Python★ 1Jul 17, 2026

MITJul 17, 2026 · metrics 1.13.0

PyPI

55Moderatehealth index

HighDimensionalEconLab/nlls_gram

Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax NNX models

Python★ 1↓ 2,882/moJul 18, 2026

MITJul 18, 2026 · metrics 1.13.0

PyPI

55Moderatehealth index

ai-hypercomputer/jetstream

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).

Python★ 451Jul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

PyPI

55Moderatehealth index

google/jetstream

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).

Python★ 451Jul 18, 2026

Apache-2.0Jul 18, 2026 · metrics 1.13.0

#jax