Calibrax

Validated against: scikit-learn and SciPy references for representative regression, classification, distance, and divergence metrics.

Documentation - Issues - Contributing

Research preview. The API will change while we iterate toward v1.0, so pin a version if you need stability. Calibrax depends on one other Avitai package, substrax, which it uses for device detection, so it is a low-commitment way to try one piece.

This is public this early on purpose. Issues, questions and pull requests genuinely steer what gets built next, and a star tells us which layer to push on.

Calibrax (Calibrate + JAX) is a unified benchmarking and metrics framework for the JAX scientific ML ecosystem. It extracts and consolidates shared benchmarking, profiling, statistical analysis, and evaluation functionality from Datarax, Artifex, and Opifex.

Features

Metrics (140 registered Tier 0 metrics, 20 domains, 4-tier architecture)

Calibrax provides a 4-tier metric system covering the full spectrum of ML evaluation. The current registry contains 140 Tier 0 pure-function metrics; Tier 1-3 APIs, optional plugins, and metric-learning losses are part of the package architecture but are not all registered metric entries today.

Functional domains: general, classification, calibration, segmentation, distance, divergence, information, ranking, statistical, clustering, fairness, forecasting, uncertainty, generative, image, text, audio, geometric, graph, manifold

Key capabilities:

• MetricRegistry with axiom-based discovery for registered Tier 0 metrics (list_true_metrics(), list_by_invariance("rotation"))
• Geometric distance hierarchy - Euclidean, Riemannian (SPD, Grassmann, Stiefel), pseudo-Riemannian (ultrahyperbolic), Finsler (Randers)
• Graph metrics - spectral distance, resistance distance, Floyd-Warshall shortest paths
• Reference checks - representative Tier 0 metrics are tested against scikit-learn and SciPy references with 1e-6 tolerance; see Peer Comparison
• Losses with masks and weights - MSE, MAE, Huber, Charbonnier, relative L2 and softmax cross-entropy take mask, weights, reduction and axis, reduced one way
• Composition - MetricCollection, WeightedMetric, MetricSuite, ThresholdMetric
• Wrappers - BootstrapMetric (confidence intervals), ClasswiseWrapper, MetricTracker, MinMaxTracker
• Metric learning losses - contrastive, triplet margin, NTXent, ArcFace, CosFace, ProxyNCA, ProxyAnchor, with hard/semi-hard negative mining

Benchmarking & Profiling

• Timing - Warm-up aware timing with JIT compilation separation
• Resource monitoring - CPU, memory, GPU memory/clock/power tracking
• Energy & carbon - Energy measurement with carbon footprint estimation
• FLOPS & roofline - XLA-level FLOP counting, roofline performance analysis
• Compilation - XLA compilation profiling and tracing
• Complexity - Algorithmic complexity analysis
• Hardware - Automatic hardware detection and capability reporting

Analysis & Infrastructure

• Statistical analysis - Bootstrap confidence intervals, hypothesis testing, effect sizes, outlier detection
• Regression detection - Direction-aware threshold checks against a stored baseline
• Comparison & ranking - Cross-configuration comparison, Pareto front analysis, aggregate scoring
• Validation - Convergence analysis and accuracy assessment
• Storage - JSON-per-run file backend with baseline management
• Exporters - W&B and MLflow integration, publication-ready LaTeX/HTML/CSV tables and matplotlib plots
• CI integration - Regression gate with git bisect automation
• Monitoring - Production alerting with configurable thresholds
• CLI - calibrax ingest|export|check|baseline|trend|summary|profile

Quick Start

import jax.numpy as jnp
from calibrax.metrics import MetricRegistry, calculate_all
from calibrax.metrics.functional.regression import mse, mae, r_squared

predictions = jnp.array([1.1, 2.3, 2.8, 4.2, 4.7])
targets = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])

# Individual metrics
print(f"MSE: {mse(predictions, targets):.4f}")
print(f"R²:  {r_squared(predictions, targets):.4f}")

# Batch computation of all registered metrics
results = calculate_all(predictions, targets, metrics=["mse", "mae", "rmse", "r_squared"])

# Registry discovery
registry = MetricRegistry()
true_metrics = registry.list_true_metrics()
rotation_inv = registry.list_by_invariance("rotation")

Installation

# Basic installation
uv pip install calibrax

# With GPU monitoring
uv pip install "calibrax[cuda12]"

# With image quality plugins (FID, Inception Score)
uv pip install "calibrax[image]"

# With text quality plugins (BERTScore)
uv pip install "calibrax[text]"

# With publication export (matplotlib)
uv pip install "calibrax[publication]"

Architecture

src/calibrax/
├── core/          Data models, protocols, adapters, result container, registry
├── profiling/     Timing, resources, GPU, energy, FLOPS, roofline, compilation,
│                  complexity, hardware, tracing, carbon
├── statistics/    Statistical analyzer, significance testing
├── analysis/      Regression, comparison, ranking, scaling, Pareto, changepoint
├── validation/    Convergence, accuracy, validation framework
├── monitoring/    Alerts, production monitoring
├── storage/       JSON store, baselines
├── exporters/     W&B, MLflow, publication-ready output
├── metrics/
│   ├── functional/   140 Tier 0 pure functions across 20 domains
│   ├── stateful/     Tier 1-2 base classes (FrozenBackboneMetric, LearnedMetric)
│   ├── learning/     Tier 3 metric learning losses and miners
│   ├── plugins/      Optional-dependency metrics (FID, BERTScore, LPIPS)
│   ├── composition.py   MetricCollection, WeightedMetric, MetricSuite, ThresholdMetric
│   ├── wrappers.py      BootstrapMetric, ClasswiseWrapper, MetricTracker, MinMaxTracker
│   └── _registry.py     MetricRegistry singleton with axiom-based discovery
├── ci/            CI regression gate, bisection engine
└── cli/           Command-line interface

Examples

Runnable examples are in examples/metrics/, available as both Python scripts and Jupyter notebooks:

Contributing

Development setup, the setup.sh flags, and the verification commands are in CONTRIBUTING.md; the contributor documentation starts at docs/contributing.

License

MIT