Metadata-Version: 2.4
Name: calibrax
Version: 0.1.0
Summary: Calibrax: Unified benchmarking framework for the JAX scientific ML ecosystem
Project-URL: Documentation, https://calibrax.readthedocs.io
Project-URL: Repository, https://github.com/avitai/calibrax
Author-email: Mahdi Shafiei <mahdi@avitai.bio>
License: MIT License
        
        Copyright (c) 2026 Mahdi Shafiei
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: benchmarking,jax,machine-learning,profiling,scientific-computing,statistics
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Natural Language :: English
Classifier: Operating System :: MacOS
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.11
Requires-Dist: click>=8.0
Requires-Dist: flax>=0.12.0
Requires-Dist: jax>=0.4.0
Requires-Dist: jaxlib>=0.4.0
Requires-Dist: jaxtyping>=0.2.20
Requires-Dist: numpy<2.4.0,>=1.24
Requires-Dist: psutil>=5.9.5
Requires-Dist: typing-extensions>=4.5
Provides-Extra: all
Requires-Dist: beartype>=0.14.1; extra == 'all'
Requires-Dist: build>=1.0.3; extra == 'all'
Requires-Dist: codecarbon>=2.0; extra == 'all'
Requires-Dist: coverage>=7; extra == 'all'
Requires-Dist: griffe>=1.7.3; extra == 'all'
Requires-Dist: hippogriffe>=0.2; extra == 'all'
Requires-Dist: interrogate>=1.5; extra == 'all'
Requires-Dist: ipykernel>=6.29.5; extra == 'all'
Requires-Dist: jax-cuda12-pjrt>=0.7.0; extra == 'all'
Requires-Dist: jax-cuda12-plugin>=0.7.0; extra == 'all'
Requires-Dist: jupytext>=1.16.6; extra == 'all'
Requires-Dist: matplotlib>=3.8; extra == 'all'
Requires-Dist: mkdocs-include-exclude-files>=0.1; extra == 'all'
Requires-Dist: mkdocs-ipynb>=0.1; extra == 'all'
Requires-Dist: mkdocs-jupyter>=0.24; extra == 'all'
Requires-Dist: mkdocs-material>=9.6.7; extra == 'all'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all'
Requires-Dist: mlflow>=2.0; extra == 'all'
Requires-Dist: nvidia-cublas-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-cuda-cupti-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-cuda-nvcc-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-cuda-nvrtc-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-cuda-runtime-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-cudnn-cu12>=9.1; extra == 'all'
Requires-Dist: nvidia-cufft-cu12>=11.4; extra == 'all'
Requires-Dist: nvidia-curand-cu12>=10.3; extra == 'all'
Requires-Dist: nvidia-cusolver-cu12>=11.7; extra == 'all'
Requires-Dist: nvidia-cusparse-cu12>=12.3; extra == 'all'
Requires-Dist: nvidia-nccl-cu12>=2.20; extra == 'all'
Requires-Dist: nvidia-nvjitlink-cu12>=12.6; extra == 'all'
Requires-Dist: nvidia-nvtx-cu12>=12.6; extra == 'all'
Requires-Dist: pre-commit>=4.2; extra == 'all'
Requires-Dist: pydeps>=1.12.17; extra == 'all'
Requires-Dist: pymdown-extensions>=10.14.3; extra == 'all'
Requires-Dist: pynvml>=11.5.0; extra == 'all'
Requires-Dist: pyright>=1.1.336; extra == 'all'
Requires-Dist: pytest-asyncio>=0.23; extra == 'all'
Requires-Dist: pytest-benchmark>=4; extra == 'all'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all'
Requires-Dist: pytest-env>=1.0.1; extra == 'all'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all'
Requires-Dist: pytest-randomly>=3.15; extra == 'all'
Requires-Dist: pytest-timeout>=2.1; extra == 'all'
Requires-Dist: pytest-xdist>=3.6; extra == 'all'
Requires-Dist: pytest>=8.3.5; extra == 'all'
Requires-Dist: radon>=6.0; extra == 'all'
Requires-Dist: ruff>=0.1.5; extra == 'all'
Requires-Dist: ruptures>=1.1; extra == 'all'
Requires-Dist: scipy>=1.10; extra == 'all'
Requires-Dist: twine>=4.0.2; extra == 'all'
Requires-Dist: wandb>=0.18; extra == 'all'
Provides-Extra: all-cpu
Requires-Dist: beartype>=0.14.1; extra == 'all-cpu'
Requires-Dist: build>=1.0.3; extra == 'all-cpu'
Requires-Dist: codecarbon>=2.0; extra == 'all-cpu'
Requires-Dist: coverage>=7; extra == 'all-cpu'
Requires-Dist: griffe>=1.7.3; extra == 'all-cpu'
Requires-Dist: hippogriffe>=0.2; extra == 'all-cpu'
Requires-Dist: interrogate>=1.5; extra == 'all-cpu'
Requires-Dist: ipykernel>=6.29.5; extra == 'all-cpu'
Requires-Dist: jupytext>=1.16.6; extra == 'all-cpu'
Requires-Dist: matplotlib>=3.8; extra == 'all-cpu'
Requires-Dist: mkdocs-include-exclude-files>=0.1; extra == 'all-cpu'
Requires-Dist: mkdocs-ipynb>=0.1; extra == 'all-cpu'
Requires-Dist: mkdocs-jupyter>=0.24; extra == 'all-cpu'
Requires-Dist: mkdocs-material>=9.6.7; extra == 'all-cpu'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all-cpu'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all-cpu'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all-cpu'
Requires-Dist: mlflow>=2.0; extra == 'all-cpu'
Requires-Dist: pre-commit>=4.2; extra == 'all-cpu'
Requires-Dist: pydeps>=1.12.17; extra == 'all-cpu'
Requires-Dist: pymdown-extensions>=10.14.3; extra == 'all-cpu'
Requires-Dist: pyright>=1.1.336; extra == 'all-cpu'
Requires-Dist: pytest-asyncio>=0.23; extra == 'all-cpu'
Requires-Dist: pytest-benchmark>=4; extra == 'all-cpu'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all-cpu'
Requires-Dist: pytest-env>=1.0.1; extra == 'all-cpu'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all-cpu'
Requires-Dist: pytest-randomly>=3.15; extra == 'all-cpu'
Requires-Dist: pytest-timeout>=2.1; extra == 'all-cpu'
Requires-Dist: pytest-xdist>=3.6; extra == 'all-cpu'
Requires-Dist: pytest>=8.3.5; extra == 'all-cpu'
Requires-Dist: radon>=6.0; extra == 'all-cpu'
Requires-Dist: ruff>=0.1.5; extra == 'all-cpu'
Requires-Dist: ruptures>=1.1; extra == 'all-cpu'
Requires-Dist: scipy>=1.10; extra == 'all-cpu'
Requires-Dist: twine>=4.0.2; extra == 'all-cpu'
Requires-Dist: wandb>=0.18; extra == 'all-cpu'
Provides-Extra: all-macos
Requires-Dist: beartype>=0.14.1; extra == 'all-macos'
Requires-Dist: build>=1.0.3; extra == 'all-macos'
Requires-Dist: codecarbon>=2.0; extra == 'all-macos'
Requires-Dist: coverage>=7; extra == 'all-macos'
Requires-Dist: griffe>=1.7.3; extra == 'all-macos'
Requires-Dist: hippogriffe>=0.2; extra == 'all-macos'
Requires-Dist: interrogate>=1.5; extra == 'all-macos'
Requires-Dist: ipykernel>=6.29.5; extra == 'all-macos'
Requires-Dist: jax-metal>=0.1.0; extra == 'all-macos'
Requires-Dist: jupytext>=1.16.6; extra == 'all-macos'
Requires-Dist: matplotlib>=3.8; extra == 'all-macos'
Requires-Dist: mkdocs-include-exclude-files>=0.1; extra == 'all-macos'
Requires-Dist: mkdocs-ipynb>=0.1; extra == 'all-macos'
Requires-Dist: mkdocs-jupyter>=0.24; extra == 'all-macos'
Requires-Dist: mkdocs-material>=9.6.7; extra == 'all-macos'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all-macos'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all-macos'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all-macos'
Requires-Dist: mlflow>=2.0; extra == 'all-macos'
Requires-Dist: pre-commit>=4.2; extra == 'all-macos'
Requires-Dist: pydeps>=1.12.17; extra == 'all-macos'
Requires-Dist: pymdown-extensions>=10.14.3; extra == 'all-macos'
Requires-Dist: pyright>=1.1.336; extra == 'all-macos'
Requires-Dist: pytest-asyncio>=0.23; extra == 'all-macos'
Requires-Dist: pytest-benchmark>=4; extra == 'all-macos'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all-macos'
Requires-Dist: pytest-env>=1.0.1; extra == 'all-macos'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all-macos'
Requires-Dist: pytest-randomly>=3.15; extra == 'all-macos'
Requires-Dist: pytest-timeout>=2.1; extra == 'all-macos'
Requires-Dist: pytest-xdist>=3.6; extra == 'all-macos'
Requires-Dist: pytest>=8.3.5; extra == 'all-macos'
Requires-Dist: radon>=6.0; extra == 'all-macos'
Requires-Dist: ruff>=0.1.5; extra == 'all-macos'
Requires-Dist: ruptures>=1.1; extra == 'all-macos'
Requires-Dist: scipy>=1.10; extra == 'all-macos'
Requires-Dist: twine>=4.0.2; extra == 'all-macos'
Requires-Dist: wandb>=0.18; extra == 'all-macos'
Provides-Extra: changepoint
Requires-Dist: ruptures>=1.1; extra == 'changepoint'
Provides-Extra: codecarbon
Requires-Dist: codecarbon>=2.0; extra == 'codecarbon'
Provides-Extra: cuda-dev
Requires-Dist: build>=1.0.3; extra == 'cuda-dev'
Requires-Dist: coverage>=7; extra == 'cuda-dev'
Requires-Dist: interrogate>=1.5; extra == 'cuda-dev'
Requires-Dist: ipykernel>=6.29.5; extra == 'cuda-dev'
Requires-Dist: jax-cuda12-pjrt>=0.7.0; extra == 'cuda-dev'
Requires-Dist: jax-cuda12-plugin>=0.7.0; extra == 'cuda-dev'
Requires-Dist: jupytext>=1.16.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cublas-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cuda-cupti-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cuda-nvcc-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cuda-nvrtc-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cuda-runtime-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-cudnn-cu12>=9.1; extra == 'cuda-dev'
Requires-Dist: nvidia-cufft-cu12>=11.4; extra == 'cuda-dev'
Requires-Dist: nvidia-curand-cu12>=10.3; extra == 'cuda-dev'
Requires-Dist: nvidia-cusolver-cu12>=11.7; extra == 'cuda-dev'
Requires-Dist: nvidia-cusparse-cu12>=12.3; extra == 'cuda-dev'
Requires-Dist: nvidia-nccl-cu12>=2.20; extra == 'cuda-dev'
Requires-Dist: nvidia-nvjitlink-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: nvidia-nvtx-cu12>=12.6; extra == 'cuda-dev'
Requires-Dist: pre-commit>=4.2; extra == 'cuda-dev'
Requires-Dist: pydeps>=1.12.17; extra == 'cuda-dev'
Requires-Dist: pynvml>=11.5.0; extra == 'cuda-dev'
Requires-Dist: pyright>=1.1.336; extra == 'cuda-dev'
Requires-Dist: pytest-asyncio>=0.23; extra == 'cuda-dev'
Requires-Dist: pytest-benchmark>=4; extra == 'cuda-dev'
Requires-Dist: pytest-cov>=6.1.1; extra == 'cuda-dev'
Requires-Dist: pytest-env>=1.0.1; extra == 'cuda-dev'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'cuda-dev'
Requires-Dist: pytest-randomly>=3.15; extra == 'cuda-dev'
Requires-Dist: pytest-timeout>=2.1; extra == 'cuda-dev'
Requires-Dist: pytest-xdist>=3.6; extra == 'cuda-dev'
Requires-Dist: pytest>=8.3.5; extra == 'cuda-dev'
Requires-Dist: radon>=6.0; extra == 'cuda-dev'
Requires-Dist: ruff>=0.1.5; extra == 'cuda-dev'
Requires-Dist: twine>=4.0.2; extra == 'cuda-dev'
Provides-Extra: dev
Requires-Dist: build>=1.0.3; extra == 'dev'
Requires-Dist: coverage>=7; extra == 'dev'
Requires-Dist: interrogate>=1.5; extra == 'dev'
Requires-Dist: ipykernel>=6.29.5; extra == 'dev'
Requires-Dist: jupytext>=1.16.6; extra == 'dev'
Requires-Dist: pre-commit>=4.2; extra == 'dev'
Requires-Dist: pydeps>=1.12.17; extra == 'dev'
Requires-Dist: pyright>=1.1.336; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
Requires-Dist: pytest-benchmark>=4; extra == 'dev'
Requires-Dist: pytest-cov>=6.1.1; extra == 'dev'
Requires-Dist: pytest-env>=1.0.1; extra == 'dev'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'dev'
Requires-Dist: pytest-randomly>=3.15; extra == 'dev'
Requires-Dist: pytest-timeout>=2.1; extra == 'dev'
Requires-Dist: pytest-xdist>=3.6; extra == 'dev'
Requires-Dist: pytest>=8.3.5; extra == 'dev'
Requires-Dist: radon>=6.0; extra == 'dev'
Requires-Dist: ruff>=0.1.5; extra == 'dev'
Requires-Dist: twine>=4.0.2; extra == 'dev'
Provides-Extra: docs
Requires-Dist: griffe>=1.7.3; extra == 'docs'
Requires-Dist: hippogriffe>=0.2; extra == 'docs'
Requires-Dist: mkdocs-include-exclude-files>=0.1; extra == 'docs'
Requires-Dist: mkdocs-ipynb>=0.1; extra == 'docs'
Requires-Dist: mkdocs-jupyter>=0.24; extra == 'docs'
Requires-Dist: mkdocs-material>=9.6.7; extra == 'docs'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'docs'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'docs'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'docs'
Requires-Dist: pymdown-extensions>=10.14.3; extra == 'docs'
Provides-Extra: gpu
Requires-Dist: jax-cuda12-pjrt>=0.7.0; extra == 'gpu'
Requires-Dist: jax-cuda12-plugin>=0.7.0; extra == 'gpu'
Requires-Dist: nvidia-cublas-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-cuda-cupti-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-cuda-nvcc-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-cuda-nvrtc-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-cuda-runtime-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-cudnn-cu12>=9.1; extra == 'gpu'
Requires-Dist: nvidia-cufft-cu12>=11.4; extra == 'gpu'
Requires-Dist: nvidia-curand-cu12>=10.3; extra == 'gpu'
Requires-Dist: nvidia-cusolver-cu12>=11.7; extra == 'gpu'
Requires-Dist: nvidia-cusparse-cu12>=12.3; extra == 'gpu'
Requires-Dist: nvidia-nccl-cu12>=2.20; extra == 'gpu'
Requires-Dist: nvidia-nvjitlink-cu12>=12.6; extra == 'gpu'
Requires-Dist: nvidia-nvtx-cu12>=12.6; extra == 'gpu'
Requires-Dist: pynvml>=11.5.0; extra == 'gpu'
Provides-Extra: image
Provides-Extra: metal
Requires-Dist: jax-metal>=0.1.0; extra == 'metal'
Provides-Extra: mlflow
Requires-Dist: mlflow>=2.0; extra == 'mlflow'
Provides-Extra: publication
Requires-Dist: matplotlib>=3.8; extra == 'publication'
Provides-Extra: scientific
Provides-Extra: stats
Requires-Dist: scipy>=1.10; extra == 'stats'
Provides-Extra: test
Requires-Dist: beartype>=0.14.1; extra == 'test'
Requires-Dist: codecarbon>=2.0; extra == 'test'
Requires-Dist: coverage>=7; extra == 'test'
Requires-Dist: matplotlib>=3.8; extra == 'test'
Requires-Dist: mlflow>=2.0; extra == 'test'
Requires-Dist: pytest-asyncio>=0.23; extra == 'test'
Requires-Dist: pytest-benchmark>=4; extra == 'test'
Requires-Dist: pytest-cov>=6.1.1; extra == 'test'
Requires-Dist: pytest-env>=1.0.1; extra == 'test'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'test'
Requires-Dist: pytest-timeout>=2.1; extra == 'test'
Requires-Dist: pytest-xdist>=3.6; extra == 'test'
Requires-Dist: pytest>=8.3.5; extra == 'test'
Requires-Dist: ruptures>=1.1; extra == 'test'
Requires-Dist: wandb>=0.18; extra == 'test'
Provides-Extra: text
Provides-Extra: wandb
Requires-Dist: wandb>=0.18; extra == 'wandb'
Description-Content-Type: text/markdown

# Calibrax

[![CI](https://github.com/avitai/calibrax/actions/workflows/ci.yml/badge.svg)](https://github.com/avitai/calibrax/actions/workflows/ci.yml)
[![Build](https://github.com/avitai/calibrax/actions/workflows/build-verification.yml/badge.svg)](https://github.com/avitai/calibrax/actions/workflows/build-verification.yml)
[![Quality](https://github.com/avitai/calibrax/actions/workflows/quality-checks.yml/badge.svg)](https://github.com/avitai/calibrax/actions/workflows/quality-checks.yml)
[![Security](https://github.com/avitai/calibrax/actions/workflows/security.yml/badge.svg)](https://github.com/avitai/calibrax/actions/workflows/security.yml)
[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
[![JAX](https://img.shields.io/badge/JAX-0.4+-green.svg)](https://github.com/jax-ml/jax)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

---

> **Early Development** — API is unstable and subject to breaking changes.
> Pin to specific commits if stability is required.

---

**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](https://github.com/avitai/datarax),
[Artifex](https://github.com/avitai/artifex), and
[Opifex](https://github.com/avitai/opifex).

## Features

### Metrics (110+ registered, 17 domains, 4 tiers)

Calibrax provides a 4-tier metric system covering the full spectrum of ML evaluation:

| Tier | Name | Pattern | Examples |
|------|------|---------|----------|
| 0 | Pure Functions | `fn(predictions, targets) -> scalar` | MSE, cosine distance, BLEU |
| 1 | Frozen Backbone | `update() -> compute() -> reset()` | FID, BERTScore, Inception Score |
| 2 | Learned | `nnx.Module` with trainable weights | LPIPS |
| 3 | Metric Learning | Differentiable embedding loss | Contrastive, Triplet, ArcFace |

**Functional domains:** regression, classification, calibration, segmentation, distance, divergence, information, ranking, statistical, clustering, fairness, image, text, audio, geometric, graph, manifold

**Key capabilities:**

- **MetricRegistry** with axiom-based discovery (`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
- **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 detection with configurable severity levels
- **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

```python
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

```bash
# Basic installation
uv pip install calibrax

# With statistical analysis (scipy)
uv pip install "calibrax[stats]"

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

# 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]"
```

## Development Setup

The recommended way to set up a development environment is with the included `setup.sh` script. It auto-detects your platform (Linux CUDA, macOS Intel, Apple Silicon), creates a virtual environment, installs all dependencies, and generates an activation script.

```bash
git clone https://github.com/avitai/calibrax.git
cd calibrax

# Standard setup with automatic GPU detection
./setup.sh

# Activate the environment
source ./activate.sh
```

### setup.sh Options

| Flag | Description |
|------|-------------|
| `--cpu-only` | Force CPU-only setup, skip GPU/Metal detection |
| `--metal` | Enable Metal acceleration on Apple Silicon Macs |
| `--deep-clean` | Clear JAX cache, pip cache, pytest cache, and other artifacts |
| `--force` | Force reinstallation even if environment exists |
| `--verbose`, `-v` | Show detailed output during setup |

```bash
# Examples
./setup.sh --cpu-only         # CPU-only development
./setup.sh --metal            # Apple Silicon with Metal
./setup.sh --force --verbose  # Force reinstall with full output
./setup.sh --deep-clean       # Clean everything and start fresh
```

### Manual Setup

If you prefer to set up manually:

```bash
git clone https://github.com/avitai/calibrax.git
cd calibrax
uv venv
uv pip install -e ".[dev,test,stats]"
uv run pre-commit install
```

## 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/   110+ Tier 0 pure functions across 17 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:

| Example | Level | Topics |
|---------|-------|--------|
| [01_quickstart.py](examples/metrics/01_quickstart.py) | Beginner | Individual metrics, `calculate_all`, registry queries |
| [02_regression_deep_dive.py](examples/metrics/02_regression_deep_dive.py) | Beginner | All 12 regression metrics, outlier sensitivity |
| [03_classification.py](examples/metrics/03_classification.py) | Intermediate | Classification, calibration, segmentation |
| [04_distances.py](examples/metrics/04_distances.py) | Intermediate | Euclidean, hyperbolic, divergences, information theory |
| [05_composition.py](examples/metrics/05_composition.py) | Intermediate | Collections, weighted metrics, quality gates, tracking |
| [06_image_quality.py](examples/metrics/06_image_quality.py) | Intermediate | PSNR, SSIM, MS-SSIM, BLEU, ROUGE |
| [07_metric_learning.py](examples/metrics/07_metric_learning.py) | Advanced | Contrastive, triplet, NTXent, ArcFace, mining |
| [08_manifold_graph.py](examples/metrics/08_manifold_graph.py) | Advanced | SPD, Grassmann, spectral distance, Floyd-Warshall |

## Development

```bash
# Run tests
uv run pytest tests/ -v --cov=calibrax --cov-report=term-missing

# Lint & format
uv run ruff check src/ tests/ --fix
uv run ruff format src/ tests/

# Type check
uv run pyright src/

# All quality checks
uv run pre-commit run --all-files

# Build documentation
uv run mkdocs build

# Convert examples to Jupyter notebooks
uv run python scripts/jupytext_converter.py batch-py-to-nb examples/metrics/
```

## License

MIT
