A differentiable, JAX-native foundation for scientific machine learning.

The Avitai stack is for any field that wants to use AI to model, predict, generate, optimize or control physical, biological and chemical systems. It is domain-agnostic by construction: differentiable data pipelines, measurement, generative modelling and scientific ML, in JAX and Flax NNX, with every stage differentiable end to end and probabilistic-first so an answer arrives with a calibrated sense of how much to trust it.

Five foundation libraries, plus two end-to-end applications in domains chosen to be as different from each other as possible: a cell, and a self-driving car. The second exists to answer honestly whether the foundation really generalises.

Avitai Bio is the company behind it, and synthetic biology is its domain: integrating multi-scale, multi-species biological data and domain knowledge into predictive cell models, so a Design-Build-Test-Learn cycle becomes faster, cheaper and more predictive by maximising what each experiment teaches rather than by running more of them. That work is the reason the stack exists and the hardest test of it; the stack itself is not specific to it.

Everything here is MIT, pre-1.0, and actively developed in the open.

┌──────────────────────────────────────────────────────────────┐
│  DiffBio                    │  DiffAV                        │
│  differentiable genomics    │  differentiable AV safety      │
├──────────────────────────────────────────────────────────────┤
│  opifex     scientific ML: PDEs, neural operators, PINNs,    │
│             quantum chemistry, UQ            (the capstone)  │
├──────────────────────────────────────────────────────────────┤
│  artifex    generative modeling: diffusion, VAE, flows,      │
│             energy-based, across images, audio, proteins     │
├──────────────────────────────────────────────────────────────┤
│  datarax    differentiable data pipelines                    │
├──────────────────────────────────────────────────────────────┤
│  calibrax   metrics, benchmarking, profiling   (measurement) │
├──────────────────────────────────────────────────────────────┤
│  substrax   training plumbing: devices, sharding, restarts   │
├──────────────────────────────────────────────────────────────┤
│              built on JAX, Flax NNX and XLA                  │
└──────────────────────────────────────────────────────────────┘

Each layer depends only on the layers below it. substrax is the bottom of the chain, depends on none of the others, and is the cleanest place to start.

Versions are the latest releases as of 2026-10-02, and every count below was measured against them rather than copied forward.

The whole foundation in one line:

pip install substrax calibrax datarax avitai-artifex opifex

• substrax is the training plumbing every other library here would otherwise write for itself: JAX process configuration, device discovery and array sharding, checkpoint save and restore, stopping rules, run logging, and job submission to remote compute. It is the only library here that depends on none of the others, which makes it the bottom of the chain.

• calibrax is the shared measurement layer, built on substrax. It registers 140 metrics across 20 domains (audio, calibration, classification, clustering, distance, divergence, fairness, forecasting, general, generative, geometric, graph, image, information, manifold, ranking, segmentation, statistical, text, uncertainty), with representative ones checked against scikit-learn and SciPy references at 1e-6. It also carries timing, GPU and energy monitoring, FLOP counting, roofline analysis and regression detection. Do not take the count on trust:

  from calibrax.metrics import MetricRegistry
  print(len(MetricRegistry().list_names()))   # 140

• datarax is a differentiable data-pipeline framework for JAX. Every stage is a Flax NNX module, so gradients flow through preprocessing and augmentation rather than stopping at the dataloader. DAG-based execution with caching and differentiable rebatching, multi-device sharding, deterministic O(1)-memory shuffling, and exact mid-epoch resume.

• artifex is a modular generative-modeling library. VAEs, GANs, diffusion, normalizing flows, energy-based, autoregressive and geometric models behind one typed interface, across images, text, audio, proteins, tabular data and time series.

• opifex is the scientific-ML capstone: neural operators (FNO, DeepONet, SFNO and more), physics-informed neural networks, equation discovery, quantum chemistry, atomistic molecular dynamics, and uncertainty quantification, probabilistic-first throughout.

Two repos take the same libraries all the way into finished domains. They exist to answer one question honestly: does the foundation actually generalise?

• DiffBio makes a genomics pipeline trainable. Hard thresholds, argmax and Smith-Waterman all block gradients, which is why bioinformatics pipelines get hand-tuned instead of learned. DiffBio replaces them with differentiable relaxations (soft quality filtering, temperature-softmax pileup, continuous Smith-Waterman), so variant calling, single-cell, RNA-seq, multi-omics and perturbation pipelines become one differentiable function you can optimize against a downstream objective. 40+ operators, 6 named pipelines.

• DiffAV makes autonomous-driving safety evaluation differentiable. A physics-constrained (bicycle-model) diffusion world-model predicts multi-agent trajectories, then RL-aligned adversarial steering pushes a scenario's adversary toward its victim while keeping the result feasible. In the measured result it moves the adversary from 7.30 m to 2.88 m from its victim while off-road feasibility improves, 0.21 to 0.08.

Being specific here is more useful than being impressive.

• Pre-1.0. APIs will change without deprecation cycles. Pin a version if you need stability.
• DiffAV is a research scaffold, not a leaderboard entry. Its trajectory minADE6 is about 5.6 m, scored with JAX-native proxies for the WOSAC metrics rather than the official implementation, so it is not comparable with leaderboard entries. The adversarial-steering behaviour is the result. It is not state of the art at trajectory prediction and we do not claim it is.
• The domain examples are worked tutorials, not hardened solutions. The foundation ships example suites spanning physics and PDEs, neural operators, quantum chemistry, molecular dynamics, vision, audio and generative modeling. They demonstrate reach; they are not production pipelines.

The foundation carries physics, chemistry, molecular dynamics, vision and generative examples out of the box, and two repos here take biology and autonomous driving all the way. If the same substrate would be useful in your field, open an issue and tell us what breaks. Early feedback genuinely steers what gets built next.

MIT, across every repository.