JAXlaxy is a curated, opinionated, and constantly updated map of the JAX ecosystem.
In 2026, JAX has evolved from a research curiosity into the backbone of global AI. It powers the world's largest Foundation Models, real-time physics simulators, and structural biology breakthroughs. While the original DeepMind "JAX Ecosystem" post (2020) is a classic archive, the universe has expanded. JAXlaxy filters the noise to show you the brightest stars.
Every entry ships with a health indicator — 🟢 Active · 🟡 Stable · 🔴 Legacy — so you can tell at a glance whether a library is the right launchpad in 2026 or a relic better admired than flown.
- ☀️ The Sun — Core & Kernels
- 🪐 The Giants — Neural Network Frameworks
- 🛰️ The Satellites — Training Infrastructure
- 🌌 Constellations — LLM & Foundation-Model Training
- 🧪 Scientific Computing & Simulation
- 📊 Probabilistic Programming
- 🤖 Reinforcement Learning & Evolution
- 🔭 Domain Libraries — Graphs, Vision, Brain Dynamics
- 🧭 The Pathfinder — What should you choose?
- 🛸 Onramps — Reference Implementations
⚠️ Legacy Radar — The warning zone
The foundational technologies that make high-performance JAX computing possible. If your stack has JAX, it has these.
- 🟢 JAX: Autograd and XLA-powered numerical computing — the sun of our universe.
- 🟢 Pallas: JAX-native kernel authoring language for writing custom TPU/GPU kernels without leaving Python.
- 🟢 OpenXLA: The open-source compiler that turns your JAX program into blazing accelerator code.
- 🟢 torchax: Run PyTorch model code directly on JAX — the bridge for teams migrating in from the PyTorch universe.
The primary solar systems for building and training neural networks in JAX.
- 🟢 Flax (NNX): The 2026 standard for DeepMind and Google-scale research. NNX brings ergonomic object-oriented state on top of JAX's functional core.
- 🟢 Equinox: Everything is a PyTree. Minimal magic, maximum transparency — the default pick for scientific ML and PyTorch converts.
- 🟢 Penzai: DeepMind's toolkit for legible, introspectable, surgically-editable neural networks — the interpretability researcher's first choice.
💡 Pragmatic multi-framework options. Teams already invested in the Keras or HuggingFace ecosystem can use Keras 3 with its JAX backend or HuggingFace Transformers Flax models — not JAX-native in design, but battle-tested bridges for real production stacks.
The utilities that keep your training loop stable, your weights safe, and your data moving at TPU speed.
- 🟢 Optax: The undisputed king of gradient processing — composable optimizer transforms used in every serious JAX training loop.
- 🟢 Orbax: The 2026 standard for multi-host checkpointing, model exporting, and resumable training at scale.
- 🟢 Grain: Deterministic, JAX-native high-throughput data loading — the TPU-era replacement for
tf.data. - 🟢 Chex: DeepMind's assertions and testing toolkit for writing JAX code you can actually trust.
- 🟢 jaxtyping: Shape-and-dtype-aware type hints that catch bugs before the first JIT compile.
- 🟡 safejax: Serialize Flax / Haiku / Equinox parameters via
safetensors— safer than pickling, portable across frameworks. - 🟢 jax-tqdm: Add a real progress bar to JIT-compiled
jax.lax.scanand training loops — one decorator, zero friction. - 🟢 JAX Toolbox: NVIDIA's nightly CI and optimized container images for running JAX workloads on H100/B200 GPUs.
- 🟢 mpi4jax: Zero-copy MPI collectives inside JIT-compiled JAX code — the bridge for classical HPC clusters.
Large-scale training stacks. Pick by workload scale and how much of the stack you want to own.
- 🟢 MaxText: Google's flagship pure-JAX LLM reference — scales from single-TPU to multi-pod without leaving Python.
- 🟢 Tunix: Google's post-training toolkit on JAX — SFT, RLHF (PPO/GRPO/DAPO), and agentic RL with tool use, built on Flax NNX.
- 🟢 Levanter: Stanford CRFM's scalable, reproducible foundation-model trainer with named tensors and bit-level determinism.
- 🟢 EasyDeL: Opinionated training and serving for Llama/Mixtral/Falcon/Qwen families in JAX — ergonomics-first.
- 🟢 kvax: A production-grade FlashAttention implementation for JAX with document-mask and context-parallel support.
- 🟡 Lorax: Automatic LoRA injection for any JAX model (Flax, Haiku, Equinox) — one line to fine-tune at a fraction of the memory.
- 🟢 FlaxDiff: Multi-node, multi-device diffusion model training on TPUs — the LLM world has a sibling.
- 🟡 EasyLM: Pre-train, fine-tune, evaluate and serve LLMs in JAX/Flax — the original reference used in several early open models; stable but slower-moving.
Where JAX truly outshines PyTorch: differentiable physics, biology, cosmology, and inverse problems.
- 🟢 AlphaFold 3: DeepMind's state-of-the-art predictor for protein, nucleic-acid and ligand structure. Weights gated; non-commercial research license.
- 🟡 jax-unirep: Fast UniRep protein-embedding models — the pragmatic starting point for sequence-level protein ML.
- 🟢 Diffrax: Numerical ODE / SDE / CDE solvers in JAX — the canonical answer for Neural ODEs and continuous-time models.
- 🟢 Optimistix: Root-finding, minimization, fixed-points and nonlinear least-squares — the "SciPy optimize" of JAX.
- 🟢 JAXopt: Hardware-accelerated, batchable, differentiable optimizers — great for bi-level and implicit-differentiation problems.
- 🟢 Brax: Massively parallel differentiable rigid-body physics — training humanoid policies on a single GPU in minutes.
- 🟢 JAX-MD: Differentiable molecular dynamics at accelerator speed — end-to-end backprop through an MD trajectory.
- 🟢 dynamiqs: High-performance, differentiable simulation of open and closed quantum systems in JAX.
- 🟢 XLB: Autodesk's differentiable, massively parallel Lattice-Boltzmann fluid solver for ML-in-the-loop CFD.
- 🟢 FDTDX: Finite-Difference Time-Domain electromagnetic simulation in JAX — design photonic devices with autograd.
- 🟢 JaxDF: Write differentiable PDE simulators with arbitrary discretizations — the building block for inverse-problem science.
- 🟢 JAX-in-Cell: Self-consistent particle-in-cell plasma simulations — classical HPC physics meets JAX autodiff.
- 🟢 foragax: Agent-based modelling framework in JAX — fast, vectorized, auto-differentiable social and ecological sims.
- 🟢 jax-cosmo: Differentiable cosmology — end-to-end gradients through large-scale-structure likelihoods.
- 🟢 astronomix: Differentiable (magneto)hydrodynamics for astrophysics — simulate galaxy-scale flows with autograd.
- 🟢 exojax: Automatically differentiable spectrum modelling of exoplanets and brown dwarfs.
- 🟢 jwave: Differentiable acoustic wave simulation — for medical-imaging and photoacoustic inverse problems.
- 🟢 SCICO: Los Alamos' scientific computational imaging — plug-and-play priors and inverse problems in JAX.
- 🟢 MBIRJAX: High-performance tomographic reconstruction — CT and 3D imaging with modern regularizers.
- 🟢 DiffeRT: Differentiable ray tracing for radio propagation — wireless-channel modelling with gradients.
- 🟢 tmmax: Vectorized transfer-matrix method for thin-film optics — the Swiss Army knife for photonic stacks.
- 🟢 vivsim: Fluid-structure interaction via the Immersed-Boundary Lattice-Boltzmann method — engineering-grade FSI with autograd.
Bayesian inference, sampling, and uncertainty — JAX's vectorized scans make MCMC fly.
- 🟢 NumPyro: The mainstream full-DSL probabilistic programming language — Pyro semantics on a JAX engine.
- 🟢 BlackJAX: Composable samplers — NUTS, HMC, SMC, VI — with no DSL lock-in. Bring your own log-density.
- 🟢 Distrax: DeepMind's lightweight distributions and bijectors library — a pragmatic alternative to TFP when you want minimum dependencies.
- 🟢 Dynamax: Probabilistic state-space models — HMMs, LGSSMs, nonlinear filters — with Kevin Murphy's seal of approval.
- 🟢 GPJax: Gaussian Processes in JAX — a didactic, extensible framework for kernel machines.
- 🟢 tinygp: The tiniest GP library — fast, elegant, and built for astronomers by Dan Foreman-Mackey.
- 🟢 flowjax: Normalizing flows built as Equinox modules — density estimation with a clean PyTree interface.
- 🟢 bayex: Bayesian optimization powered by JAX — hyperparameter tuning that runs inside your training job.
- 🟡 Oryx: Probabilistic programming via program transformations — inside TensorFlow Probability, niche but powerful for researchers.
💡 Also worth knowing:
tfp.substrates.jax— TensorFlow Probability's distributions, MCMC, and VI running on a pure JAX substrate.
End-to-end JIT'd training loops and hardware-accelerated environments — the JAX RL stack trains in minutes what TensorFlow took hours to simulate.
- 🟡 PureJaxRL: Fully vectorized, end-to-end JIT'd RL pipelines — PPO on 2048 envs without leaving JAX. Low recent commit activity, but remains the canonical reference for the JAX-native RL-loop pattern.
- 🟢 Jumanji: InstaDeep's suite of industry-driven, hardware-accelerated RL environments — from bin-packing to routing.
- 🟢 gymnax: Classic Gym environments re-implemented in JAX — CartPole, Atari-lite, bsuite, and more, all JIT-compatible.
- 🟢 Pgx: Vectorized board-game environments with an AlphaZero reference — Chess, Shogi, Go at scale.
- 🟢 NAVIX: MiniGrid reimplemented in pure JAX — RL gridworlds that train in seconds, not hours.
- 🟢 QDax: Quality-Diversity optimization — MAP-Elites and neuro-evolution at accelerator speed.
- 🟢 evosax: JAX-based evolutionary strategies — CMA-ES, OpenAI-ES, NSLC, ready to vectorize.
- 🟢 RLax: DeepMind's RL building blocks — value functions, distributional losses, exploration — the LEGO set, not the agent.
- 🟢 Mctx: DeepMind's Monte-Carlo Tree Search primitives in native JAX — MuZero-style planning, vectorized.
💡 For continuous-control physics environments, see Brax in the Scientific Computing section — it doubles as a world-class RL env suite. 💡 EvoJAX — the original "put ES on TPU" toolkit — has been archived; see the Legacy Radar. Use evosax or QDax above.
Specialized stellar systems: graphs, vision, neuroscience, and privacy-preserving compute.
- 🟡 PGMax: Discrete probabilistic graphical models with loopy-BP and smooth-minimum-sum inference in JAX.
⚠️ Jraph — the de facto GNN library in JAX — was archived by DeepMind. See the Legacy Radar; no drop-in JAX-native successor exists yet.
- 🟢 Scenic: Google Research's JAX/Flax library for vision transformers, video, and multi-modal research — the living vision codebase in JAX.
- 🟢 dm_pix: DeepMind's image-processing primitives for JAX — JIT-friendly augmentations and color ops.
⚠️ Note on vision model zoos: most Flax/Equinox pre-trained-weight repos (FlaxVision, jax-models, Eqxvision) have gone dormant. See the Legacy Radar and prefer Scenic or HuggingFace Transformers' Flax models for new work.
- 🟢 BrainPy: Computational neuroscience and brain-inspired computing — differentiable spiking networks and neural dynamics.
- 🟢 brainunit: Physical units and unit-aware arithmetic inside JAX — make your neuroscience code dimensionally safe.
- 🟢 brainstate: State-based program compilation for brain-dynamics models — augmenting JAX's functional core with stateful ergonomics.
- 🟢 dendritex: Compartmental dendritic neuron modelling in JAX — cable-equation dynamics at GPU speed.
- 🟢 Spyx: Spiking Neural Networks in JAX — neuromorphic-style learning with modern accelerators.
- 🟢 OTT-JAX: Optimal transport — Sinkhorn, low-rank Gromov-Wasserstein, and neural OT — the reference toolkit in JAX.
- 🟢 Coreax: GCHQ's coreset algorithms for compressing large datasets while preserving statistical structure.
- 🟢 SPU: A compiler + runtime for running JAX programs under Secure Multi-Party Computation — privacy-preserving ML, the compiler way.
Real 2026 user journeys → recommended stars. Pick the row that matches your mission.
| Your Mission | Recommended | Status | Why |
|---|---|---|---|
| Large-scale LLM training on TPU | MaxText | 🟢 | Pure-JAX, scales multi-pod, battle-tested on Gemini-class workloads |
| Post-training LLMs (SFT + RLHF) | Tunix | 🟢 | PPO/GRPO/DAPO with tool-using agents, on Flax NNX |
| Foundation-model research (custom arch) | Flax NNX | 🟢 | DeepMind's new OO-ergonomic API — Haiku's successor |
| Scientific ML / PyTorch-style transparency | Equinox | 🟢 | Callable PyTrees, minimal magic, strong sci-ML adoption |
| Neural ODEs / continuous-time models | Diffrax | 🟢 | The canonical differential-equation solver in JAX |
| Differentiable physics simulation | Brax or JAX-MD | 🟢 | Brax = rigid body, JAX-MD = molecular dynamics |
| Probabilistic modelling (full DSL) | NumPyro | 🟢 | Pyro-lineage, fast, mainstream |
| Sampling only (MCMC / SMC / VI) | BlackJAX | 🟢 | Composable samplers, no DSL lock-in |
| RL research | PureJaxRL + Jumanji / gymnax | 🟡 | End-to-end JIT'd loops (PureJaxRL is in maintenance but still canonical) + 🟢 accelerator-native envs |
| Protein / biomolecular structure | AlphaFold 3 | 🟢 | DeepMind canonical (research license) |
| Interpretability / model surgery | Penzai | 🟢 | DeepMind's introspective modelling library |
| TPU/GPU kernel authoring | Pallas | 🟢 | JAX-native — write kernels without leaving Python |
| Concern | Use | Why |
|---|---|---|
| Optimizers | Optax | Composable gradient transforms — universal adoption |
| Checkpointing | Orbax | Multi-host, resumable, the 2026 standard |
| Data loading | Grain | Deterministic, JAX-native, replaces tf.data |
| Testing & invariants | Chex | DeepMind's assertions library for JAX code |
| Shape-safe types | jaxtyping | Catches shape bugs before JIT |
Canonical codebases to read when you're learning how idiomatic JAX is written. Pick one near your mission and study it.
- 🚀 MaxText Examples — see how a production-grade, multi-pod LLM training loop is structured in pure JAX.
- 🧭 PureJaxRL Tutorials — the clearest demonstration of end-to-end JIT'd training-loop design in the JAX world.
- 🧪 JAX-MD Notebooks — differentiable molecular dynamics from first principles, with narrative tutorials.
- 🔬 Penzai Tutorials — model introspection and editing — a great example of composable PyTree APIs.
- 📊 NumPyro Examples — Bayesian inference recipes, from linear regression to deep GPs.
- 🌸 Flax NNX Guides — the official "how to think in NNX" walkthrough; the best starting point if you're coming from Haiku or PyTorch.
These pioneers lit the way, but the galaxy has moved on. Each entry explains what to use instead — if you're reading old tutorials, read this first.
- 🔴 Haiku (dm-haiku) — In maintenance mode. DeepMind's new research has shifted to Flax NNX. → Use Flax NNX.
- 🔴 Trax — Effectively abandoned (Google). → Use MaxText for scale, Flax NNX for research.
- 🔴 Objax — No longer actively developed. → Use Equinox (similar OO feel) or Flax NNX.
- 🔴 Elegy — Unmaintained. → Use Flax NNX or Equinox directly.
- 🔴 SymJAX — Superseded by native JAX tracing. → Use JAX itself.
- 🔴 Parallax — Archived experimental project. → Use Equinox for "immutable modules" ergonomics.
- 🔴 mcx — Sampling DSL superseded by the community. → Use BlackJAX.
- 🔴 Coax — Slowed to a crawl. → Use PureJaxRL.
- 🔴 EvoJAX — Archived by Google. → Use evosax (general ES) or QDax (Quality-Diversity).
- 🔴 Jraph — Archived by DeepMind in 2024. → No drop-in JAX-native successor. For new GNN work, consider PyTorch Geometric or compose message-passing primitives with Equinox.
- 🔴 FlaxVision / jax-models / Eqxvision — Dormant vision model zoos. → Use Scenic or HuggingFace Transformers' Flax models.
- 🟡 Pax / Praxis — Still ships, but Google's external narrative has moved to MaxText + Flax NNX. Google-internal lineage; new users should not start here.
- 🟡 Neural Tangents — Low activity. Stays here because it's the canonical library for its (niche) infinite-width-network research area.
- 🟡 FedJAX — Federated-learning niche with minimal recent activity. No clear successor in JAX; still worth knowing if you work on FL.
The JAX multiverse is expanding. If you see a new star — or see a dying one — open an Issue or a PR. See CONTRIBUTING.md for cosmic guidelines.
To audit the health of the stars yourself, run:
python scripts/health_check.pyThis fetches last-commit dates and star counts for every linked repo and flags classifications that look stale.
The galaxy is also audited automatically every Monday via GitHub Actions — dying stars get flagged in an Issue before you ever need to look. See CONTRIBUTING.md for details.
Prefer to see the galaxy? The Observatory is a cinematic 3D map of every entry — size-by-stars, color-by-health, click-to-explore. cd observatory && npm run dev from a fresh clone.
Maintained with ❤️ by the JAX community. 2026 Edition. Let's map the stars. 🌌