9 projects
calibrax
JAX benchmarking, profiling and evaluation metrics for Flax NNX: a metric registry across regression, classification, calibration, uncertainty, forecasting, generative, image, text, audio, graph and fairness domains, 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, diffusion, normalizing flows, energy-based, autoregressive and geometric models across image, text, audio, molecular, protein, tabular and time-series modalities, with a shared Trainer and evaluation through calibrax
opifex
Scientific machine learning for JAX/Flax NNX: neural operators, physics-informed networks, E(3)-equivariant atomistic potentials, differentiable DFT, SINDy equation discovery and uncertainty quantification, with PDEBench benchmarking
substrax
JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging, and job runs on Modal, SkyPilot or locally
diffbio
End-to-end differentiable bioinformatics for JAX/Flax NNX: alignment, mapping, 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 datarax, artifex, opifex and calibrax
cellifex
Cellifex — an AI/ML-native cell-state dynamics foundation model
diffav
DiffAV: Physics-informed, RL-aligned evaluation engine for autonomous driving
fluctifex
Differentiable Earth-system modeling in JAX: assimilate raw observations, forecast, and control.
