Flax Nnx Python Packages

substrax

JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging.

calibrax

JAX benchmarking, profiling and evaluation metrics for Flax NNX: a registry of pure-function metrics (regression, classification, calibration, uncertainty, forecasting, generative, image, text, audio, graph, fairness), XLA FLOP counting, roofline analysis, GPU and energy monitoring, regression detection, publication exports, W&B and MLflow.

datarax

Differentiable data pipelines for JAX/Flax NNX: HuggingFace, TFDS and ArrayRecord sources, augmentation stages, scan-based epochs, exact mid-epoch resume.

avitai-artifex

Generative modeling for JAX/Flax NNX: VAEs, GANs, DDPM/DDIM/score/DiT/latent diffusion, normalizing flows, energy-based, autoregressive and geometric (SE(3), protein, point cloud, mesh) models; image, text, audio, molecular, protein, tabular and time-series modalities; a shared Trainer with callbacks and extensions; evaluation through calibrax.

diffbio

End-to-end differentiable bioinformatics for JAX/Flax NNX: soft Smith-Waterman alignment, read mapping and assembly, variant calling, RNA-seq, single-cell, epigenomics, CRISPR, metabolomics, multi-omics, protein and RNA structure, molecular dynamics and drug-discovery operators composed into trainable pipelines on the Avitai stack.

diffav

Physics-informed, RL-aligned evaluation engine for autonomous driving: JAX/Flax NNX trajectory-diffusion world models on the Waymo Open Dataset, DPO and reward-guided adversarial steering toward the scenarios that break AV stacks, physics feasibility, occupancy flow, WOSAC-style realism metrics and sensor simulation.

flax-metrics

Flax NNX implementation of common metrics.

jaxnn

JaxNN: Foundation Models in JAX/Flax

lynnx

Deep reinforcement learning algorithms and environments, written completely in JAX + Flax NNX for end-to-end GPU-accelerated training.