Metadata-Version: 2.5
Name: calibrax
Version: 0.1.5
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.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.12
Requires-Dist: click>=8.0
Requires-Dist: flax>=0.12.9
Requires-Dist: jax>=0.11.1
Requires-Dist: jaxlib>=0.11.1
Requires-Dist: jaxtyping>=0.2.20
Requires-Dist: numpy>=2.1
Requires-Dist: psutil>=5.9.5
Requires-Dist: substrax>=0.1.0
Requires-Dist: typing-extensions>=4.5
Provides-Extra: all
Requires-Dist: bandit[toml]>=1.8.6; 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: hypothesis>=6.100; extra == 'all'
Requires-Dist: import-linter>=2.0; extra == 'all'
Requires-Dist: interrogate>=1.5; extra == 'all'
Requires-Dist: ipykernel>=6.29.5; extra == 'all'
Requires-Dist: jax[cuda12]>=0.11.1; 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: pip-audit>=2.7; extra == 'all'
Requires-Dist: pre-commit>=4.2; extra == 'all'
Requires-Dist: pydeps>=1.12.17; extra == 'all'
Requires-Dist: pydoclint>=0.5; extra == 'all'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all'
Requires-Dist: pymdown-extensions>=10.16.1; 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-codspeed>=4.4.0; 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>=9.0.3; 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: scikit-learn>=1.4; 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: bandit[toml]>=1.8.6; 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: hypothesis>=6.100; extra == 'all-cpu'
Requires-Dist: import-linter>=2.0; 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: pip-audit>=2.7; extra == 'all-cpu'
Requires-Dist: pre-commit>=4.2; extra == 'all-cpu'
Requires-Dist: pydeps>=1.12.17; extra == 'all-cpu'
Requires-Dist: pydoclint>=0.5; extra == 'all-cpu'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all-cpu'
Requires-Dist: pymdown-extensions>=10.16.1; 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-codspeed>=4.4.0; 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>=9.0.3; 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: scikit-learn>=1.4; 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: bandit[toml]>=1.8.6; 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: hypothesis>=6.100; extra == 'all-macos'
Requires-Dist: import-linter>=2.0; 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: pip-audit>=2.7; extra == 'all-macos'
Requires-Dist: pre-commit>=4.2; extra == 'all-macos'
Requires-Dist: pydeps>=1.12.17; extra == 'all-macos'
Requires-Dist: pydoclint>=0.5; extra == 'all-macos'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all-macos'
Requires-Dist: pymdown-extensions>=10.16.1; 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-codspeed>=4.4.0; 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>=9.0.3; 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: scikit-learn>=1.4; 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: bandit[toml]>=1.8.6; extra == 'cuda-dev'
Requires-Dist: build>=1.0.3; extra == 'cuda-dev'
Requires-Dist: coverage>=7; extra == 'cuda-dev'
Requires-Dist: import-linter>=2.0; extra == 'cuda-dev'
Requires-Dist: interrogate>=1.5; extra == 'cuda-dev'
Requires-Dist: ipykernel>=6.29.5; extra == 'cuda-dev'
Requires-Dist: jax[cuda12]>=0.11.1; 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: pip-audit>=2.7; extra == 'cuda-dev'
Requires-Dist: pre-commit>=4.2; extra == 'cuda-dev'
Requires-Dist: pydeps>=1.12.17; extra == 'cuda-dev'
Requires-Dist: pydoclint>=0.5; extra == 'cuda-dev'
Requires-Dist: pygments<3.0,>=2.20; 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>=9.0.3; 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: cuda12
Requires-Dist: jax[cuda12]>=0.11.1; extra == 'cuda12'
Requires-Dist: nvidia-cublas-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-cuda-cupti-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-cuda-nvcc-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-cuda-nvrtc-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-cuda-runtime-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-cudnn-cu12>=9.1; extra == 'cuda12'
Requires-Dist: nvidia-cufft-cu12>=11.4; extra == 'cuda12'
Requires-Dist: nvidia-curand-cu12>=10.3; extra == 'cuda12'
Requires-Dist: nvidia-cusolver-cu12>=11.7; extra == 'cuda12'
Requires-Dist: nvidia-cusparse-cu12>=12.3; extra == 'cuda12'
Requires-Dist: nvidia-nccl-cu12>=2.20; extra == 'cuda12'
Requires-Dist: nvidia-nvjitlink-cu12>=12.6; extra == 'cuda12'
Requires-Dist: nvidia-nvtx-cu12>=12.6; extra == 'cuda12'
Requires-Dist: pynvml>=11.5.0; extra == 'cuda12'
Provides-Extra: dev
Requires-Dist: bandit[toml]>=1.8.6; extra == 'dev'
Requires-Dist: build>=1.0.3; extra == 'dev'
Requires-Dist: coverage>=7; extra == 'dev'
Requires-Dist: import-linter>=2.0; 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: pip-audit>=2.7; extra == 'dev'
Requires-Dist: pre-commit>=4.2; extra == 'dev'
Requires-Dist: pydeps>=1.12.17; extra == 'dev'
Requires-Dist: pydoclint>=0.5; extra == 'dev'
Requires-Dist: pygments<3.0,>=2.20; 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>=9.0.3; 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: pygments<3.0,>=2.20; extra == 'docs'
Requires-Dist: pymdown-extensions>=10.16.1; extra == 'docs'
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: performance
Requires-Dist: pytest-codspeed>=4.4.0; extra == 'performance'
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: hypothesis>=6.100; 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>=9.0.3; extra == 'test'
Requires-Dist: ruptures>=1.1; extra == 'test'
Requires-Dist: scikit-learn>=1.4; 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.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
[![JAX](https://img.shields.io/badge/JAX-0.11+-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)

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

[Documentation](https://calibrax.readthedocs.io/en/latest/) - [Issues](https://github.com/avitai/calibrax/issues) - [Contributing](CONTRIBUTING.md)

---

> **Research preview.** The API will change while we iterate toward v1.0, so pin a version if you
> need stability. Calibrax is the most standalone library in the Avitai stack: it depends on none
> of the others, 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](https://github.com/avitai/datarax),
[Artifex](https://github.com/avitai/artifex), and
[Opifex](https://github.com/avitai/opifex).

## Features

### Metrics (137 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 137 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.

| 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 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](docs/user-guide/peer-comparison.md)
- **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[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/   137 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:

| 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 | Same-shape 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 |

## Contributing

Development setup, the `setup.sh` flags, and the verification commands are in
[CONTRIBUTING.md](CONTRIBUTING.md); the contributor documentation starts at
[docs/contributing](docs/contributing/index.md).

## License

MIT
