A JAX/Flax NNX scientific ML stack — from data pipelines and evaluation primitives up through differentiable bioinformatics.

Avitai Bio develops a layered, open-source ecosystem for scientific machine learning, built natively on JAX and Flax NNX. Each layer is a standalone library; together they compose into end-to-end differentiable pipelines for the life sciences.

┌──────────────────────────────────────────────────────┐
│  DiffBio   — differentiable bioinformatics pipelines │
├──────────────────────────────────────────────────────┤
│  Artifex   │  Opifex                                 │
│  generative│  scientific ML                          │
│  models    │  (PINNs, neural operators, discovery)   │
├──────────────────────────────────────────────────────┤
│  Datarax   — differentiable data pipelines           │
│  Calibrax  — benchmarking, profiling, metrics        │
├──────────────────────────────────────────────────────┤
│              JAX  ·  Flax NNX  ·  XLA                │
└──────────────────────────────────────────────────────┘

• datarax — A differentiable data pipeline framework for JAX. DAG-based execution, deterministic shuffling, multi-device sharding, and gradients that flow through every operator. Built so preprocessing, augmentation, and even synthesis can be optimized end-to-end.
• calibrax — Unified benchmarking, profiling, and metrics for the JAX scientific ML ecosystem. 110+ registered metrics across 17 domains, plus timing, GPU/energy monitoring, FLOP counting, roofline analysis, regression detection, and CI integration.

• artifex — A research-focused modular generative modeling library. VAEs, GANs, diffusion, normalizing flows, energy-based, autoregressive, and geometric models, with unified interfaces across images, text, audio, proteins, tabular, and time series.
• opifex — A unified scientific machine learning framework. Neural operators (FNO, DeepONet, SFNO, and 20+ more), physics-informed neural networks, equation discovery (SINDy), uncertainty quantification, and quantum chemistry — all probabilistic-first and JAX-native.

• DiffBio — End-to-end differentiable bioinformatics pipelines. 40+ differentiable operators (alignment, variant calling, single-cell, epigenomics, RNA-seq, drug discovery, protein/RNA structure) composed into trainable pipelines. Replaces hard thresholds and argmax decisions with continuous relaxations so entire workflows can be optimized by gradient descent.

Each layer is independently useful — you can use Datarax without ever touching biology, or Calibrax to benchmark unrelated JAX projects. But the layers are designed to compose: DiffBio uses Datarax's operator contracts, Artifex's modeling components, Opifex's optimization primitives, and Calibrax's evaluation harness. The goal is a coherent, type-checked, JIT-compiled path from raw data through model to peer-reviewed metric — without leaving JAX.

Most projects are in early, active development with unstable APIs. They are suitable for research and experimentation today; pin to specific commits if you need stability. See individual repository READMEs for current status, install instructions, and quickstarts.

All projects are released under the MIT License.