Metadata-Version: 2.5
Name: xlaxiom
Version: 3.0.0
Summary: Universal numerical verification, kernel fuzzing, and differential testing for PyTorch, JAX, Triton, and OpenXLA.
Project-URL: Homepage, https://github.com/Hardonian/XLAxiom
Project-URL: Documentation, https://github.com/Hardonian/XLAxiom#readme
Project-URL: Repository, https://github.com/Hardonian/XLAxiom
Project-URL: Issues, https://github.com/Hardonian/XLAxiom/issues
Project-URL: Changelog, https://github.com/Hardonian/XLAxiom/blob/main/CHANGELOG.md
Author: XLAxiom contributors
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: assert-close,compiler-testing,differential-testing,failure-minimization,fp8,jax,kernel-fuzzing,numerical-testing,openxla,pallas,property-based-testing,pytorch,quantization,reproducibility,triton,xla
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Testing
Classifier: Typing :: Typed
Requires-Python: <3.14,>=3.11
Requires-Dist: numpy>=1.24
Provides-Extra: all
Requires-Dist: jax>=0.4.30; extra == 'all'
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]<0.11,>=0.5; (platform_system == 'Linux') and extra == 'cuda12'
Provides-Extra: cuda13
Requires-Dist: jax[cuda13]<0.11,>=0.5; (platform_system == 'Linux') and extra == 'cuda13'
Provides-Extra: dev
Requires-Dist: build>=1.2; extra == 'dev'
Requires-Dist: jax>=0.4.30; extra == 'dev'
Requires-Dist: mypy>=1.13; extra == 'dev'
Requires-Dist: pytest>=8.3; extra == 'dev'
Requires-Dist: ruff>=0.8; extra == 'dev'
Requires-Dist: twine>=6.1; extra == 'dev'
Provides-Extra: jax
Requires-Dist: jax>=0.4.30; extra == 'jax'
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == 'torch'
Description-Content-Type: text/markdown

# XLAxiom — turn numerical discrepancies into evidence

[![CI](https://github.com/Hardonian/XLAxiom/actions/workflows/ci.yml/badge.svg)](https://github.com/Hardonian/XLAxiom/actions/workflows/ci.yml)
[![License: Apache-2.0](https://img.shields.io/badge/License-Apache--2.0-blue.svg)](LICENSE)
[![Python 3.11–3.13](https://img.shields.io/badge/python-3.11%E2%80%933.13-blue.svg)](https://www.python.org/)

**XLAxiom turns a suspicious JAX/OpenXLA numerical result into a minimized, reproducible, portable, evidence-backed case.** It gives compiler, platform, and ML infrastructure teams a disciplined path from “this changed” to an artifact another engineer can inspect, rerun, and compare.

It is independently developed and is not affiliated with, sponsored by, or endorsed by Google.

> **Portfolio signal:** XLAxiom is the ML-infrastructure flagship: a reproducible testing and evidence pipeline for turning JAX/OpenXLA numerical discrepancies into small, reviewable engineering cases.

## Pick your path

| If you need to… | Start here |
| --- | --- |
| Prove the pipeline works safely | [Run the five-minute synthetic control](#five-minute-confidence-check). |
| Investigate a real numerical discrepancy | Follow the [artifact handoff workflow](docs/artifacts.md). |
| Add bounded regression coverage to CI | Run the [curated property corpus](docs/corpus.md). |
| Compare CPU, GPU, or version-specific evidence | Use the [support matrix](docs/support.md) and [quickstart](docs/quickstart.md). |
| Inspect physical TPU qualification | Review the [retained 4-chip TPU evidence](evidence/tpu/manifest.json). |
| Qualify TPU workflows without accelerator cost | Run the [offline TPU qualification](docs/offline-tpu.md). |
| Integrate with Google Cloud, TPUs & Colab | Read the [Google Ecosystem guide](docs/google-ecosystem.md). |
| Install from source or a verified wheel | Read the [installation guide](docs/installation.md). |

## Everyday Python API

XLAxiom provides drop-in numerical assertions, custom kernel fuzzing, and layer parity auditing for daily ML development:

```python
import xlaxiom as xl

# 1. Drop-in assertion with root-cause diagnostics, SNR & cosine similarity
xl.assert_close(actual_tensor, reference_tensor, rtol=1e-4, atol=1e-5)

# 2. Differentially fuzz custom Triton, Pallas, or CUDA kernels across dynamic shapes
res = xl.fuzz_kernel(
    my_custom_kernel,
    reference_fn,
    shape_spec=((1, 8), (16, 128), [32, 64]),
    dtypes=["float32"],
)
# Automatically shrinks failing tensors to minimal reproducing shapes!

# 3. Audit layer-by-layer activation parity between PyTorch, JAX, or compiled models
report = xl.audit_layer_parity(model_eager, model_compiled, sample_input)
print(report.format_table())  # Pinpoints exact layer where divergence originates!

# 4. Audit quantization precision loss and outlier clipping risk
audit = xl.audit_quantization_risk(activations, target_dtype="fp8_e4m3")
print(audit.format_summary())  # Outlier clipping %, underflow %, and SNR


# 5. Continuous Integration regression gate & PR comment generator
diff_obj = xl.diff(candidate_tensor, baseline_tensor)
gate = xl.evaluate_ci_gate([("model_output", diff_obj)], min_snr_db=30.0)
print(gate.summary_markdown)  # Ready to post to GitHub PR comments or step summaries!

# 6. Checkpoint weight parity audit (Safetensors / .npz / PyTorch state_dict)
ckpt_report = xl.audit_checkpoint_parity("base_model.safetensors", "quantized_model.safetensors")
print(ckpt_report.format_summary())  # Flags missing/extra keys and top divergent weights

# 7. LLM serving engine parity verification (vLLM / SGLang / TensorRT-LLM vs Reference)
serving_report = xl.audit_serving_parity(vllm_engine, hf_reference, sample_prompts)
print(serving_report.format_summary())  # Top-1 token agreement %, Top-5 overlap, and KL divergence

# 8. Automated empirical tolerance calibration (eliminates flaky CI assertions)
tolerances = xl.calibrate_tolerances(custom_kernel, sample_input, runs=10)
print(tolerances.format_summary())  # Recommends statistically robust rtol/atol with safety headroom


# 9. Hypothesis-like metamorphic invariant testing for neural network operations
@xl.test_invariants(invariants=["batch_padding", "scale_invariance"])
def test_attention_layer(x):
    return my_attention(x)
```

## Five-minute confidence check

Prerequisite: Python 3.11–3.13 and [uv](https://docs.astral.sh/uv/). From a source checkout:

```console
uv sync --extra dev --frozen
uv run xlaxiom doctor
uv run xlaxiom generate --synthetic --seed 7 --output results/original.json
uv run xlaxiom minimize --case results/original.json --output results/reduced.json
uv run xlaxiom reproduce --case results/reduced.json
uv run xlaxiom report --case results/reduced.json --output results/report
uv run xlaxiom bundle create --case results/reduced.json --output results/control.zip
uv run xlaxiom bundle validate --bundle results/control.zip
```

This produces a small, checksummed control case, a fresh-process reproduction verdict, a JSON/Markdown report, and a validated portable bundle. `--synthetic` injects a clearly labelled offset after JAX execution: it proves the workflow, but is never evidence of a JAX, XLA, hardware, or model defect.

For a real observation, start with the [artifact guide](docs/artifacts.md). It explains the required minimization, fresh-process reproduction, environment evidence, and safe handoff boundaries before an upstream issue is considered.

## What XLAxiom does

```text
observe → classify → minimize → reproduce in a fresh process
        → bundle evidence → compare environments
```

- **Records the evidence.** It captures versions, device placement, numerical policy, outputs, and an explicit predicate instead of treating a mismatch as a diagnosis.
- **Keeps cases small.** Failure-preserving minimization and topological subgraph extraction reduce a reproducible observation without silently changing what failed.
- **Pinpoints compiler regressions.** Automated bisection binary-searches compiler releases in isolated environments to identify exact breaking changes.
- **Makes handoffs portable.** Declarative JSON workloads, copy-pasteable 10-line reproduction scripts, triage markdown packages, checksummed reports, and bundles let another engineer or upstream maintainer validate the observation immediately.
- **Fails closed on hardware.** Accelerator requests must prove actual placement; unavailable GPU/TPU hardware is never represented as a CPU success.

XLAxiom is the evidence layer between a suspicious numerical result and a focused engineering investigation. It is not a general correctness oracle, performance benchmark suite, or automatic root-cause analyzer.

## The normal workflow

1. **Capture an observation.** Compare NumPy against JAX eager/JIT or execute a reviewed relation with an explicit policy.
2. **Reduce it without changing it.** Minimize a repeatable non-synthetic predicate and preserve its reduction record.
3. **Reproduce it independently.** Re-run in a fresh process, then record a second host, device, or runtime when the claim requires it.
4. **Hand it off safely.** Generate a report and validated bundle; attach only the evidence that supports the stated scope.

The [artifact lifecycle](docs/artifacts.md) has the exact command-level process. The [issue readiness check](docs/quickstart.md#before-escalating-a-real-observation) is deliberately conservative: an observation is not a diagnosis, and missing external evidence is never inferred.

## Evidence you can trust

The repository is designed to make scope visible rather than inflate it:

- The CPU path, checksummed artifacts, failure containment, and reviewed property corpus are covered by automated tests and release verification.
- The curated corpus contains forty explicit relations across ordinary, extreme, signed-zero, non-finite, integer, mixed-dtype, and x64 inputs with verified mutation sharpness.
- Windows/Linux, selected CUDA, and 4-chip Google Cloud TPU v5e evidence are retained as concrete records; untested hardware remains capability- and evidence-gated.

See the [support and verification matrix](docs/support.md) for what is supported versus merely implemented, and the [experiment records](docs/experiments/) for the underlying evidence—not a condensed lab notebook in this README.

## Command map

| Goal | Primary commands |
| --- | --- |
| Ingest live code | `trace` |
| Inspect an environment | `doctor`, `environment capture`, `runtime compare` |
| Investigate one case | `generate`, `run`, `compare`, `minimize`, `reproduce` |
| Pinpoint compiler regression | `bisect` |
| Prepare upstream triage/handoff | `triage`, `report`, `bundle create`, `bundle validate`, `readiness` |
| Test known relations | `property list`, `property run`, `property campaign`, `relation` |
| Compare changes safely | `matrix run`, `property compare`, `historical-study` |
| Validate accelerators | `gpu doctor`, `gpu validate`, `gpu campaign` |
| Google Ecosystem & Cloud TPUs | `cloud tpu-run`, `cloud tpu-status`, `cloud vertex-spec`, `cloud xpk-spec` |
| Developer surfaces & Colab | `colab`, `xla pass-isolate` |
| Enterprise Governance & Telemetry | `bigquery schema`, `bigquery export`, `kms sign`, `kms verify` |
| Universal Parity & Quantization | `diff`, `quant-audit` |

Run `xlaxiom <command> --help` for the complete interface. The [documentation index](docs/index.md) routes each workflow to the right level of detail.

## Safety and interpretation

Workloads are deterministic JSON data selected from a closed operation registry; XLAxiom does not load pickle files or execute workload-supplied code. Numerical workloads can still consume meaningful resources, so run untrusted artifacts with suitable process and resource isolation. Read [SECURITY.md](SECURITY.md) before sharing artifacts or reporting a vulnerability.

NumPy is useful but is not an independent oracle for every JAX semantic, dtype-promotion rule, compiler setting, or hardware implementation. A mismatch directs investigation; it never establishes root cause or severity by itself.

## Development and release

```console
uv sync --extra dev --frozen
uv run pytest
uv run ruff check .
uv run ruff format --check .
uv run mypy src
uv build --out-dir .release-dist
uv run twine check .release-dist/*
uv run python scripts/verify_distribution.py --dist-dir .release-dist
```

The [release policy](docs/release.md) defines the required quality gates, reproducible-build expectations, and Trusted Publishing boundary. The [roadmap](docs/roadmap.md) distinguishes implemented capability from physical-hardware evidence that cannot be manufactured locally.
