XLAxiom — turn numerical discrepancies into evidence

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

Everyday Python API

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

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)


# 10. Bit-level ULP distance assertions across float16/float32/float64
xl.assert_ulp_close(actual_tensor, reference_tensor, max_ulp=2.0)

# 11. Record-once, verify-forever golden snapshot testing
xl.check_golden("model_golden.npz", output_tensor, max_ulp=4.0)

# 12. Serving engine top-k logprob parity (vLLM / SGLang / OpenAI-compatible API)
topk_report = xl.audit_topk_parity(vllm_topk_fn, hf_topk_fn, prompts=["Hello world"])
print(topk_report.format_summary())

Turnkey GitHub Actions CI/CD Integration

Drop automated numerical regression gates into your GitHub workflows with zero boilerplate:

- name: Verify Numerical Parity Gate
  uses: Hardonian/XLAxiom@v3
  with:
    actual: "candidate_weights.npz"
    desired: "reference_weights.npz"
    rtol: "1e-4"
    atol: "1e-5"
    min-snr: "25.0"
    comment-on-pr: "true"

Five-minute confidence check

Prerequisite: Python 3.11–3.13 and uv. From a source checkout:

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. It explains the required minimization, fresh-process reproduction, environment evidence, and safe handoff boundaries before an upstream issue is considered.

What XLAxiom does

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 has the exact command-level process. The issue readiness check 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 for what is supported versus merely implemented, and the experiment records for the underlying evidence—not a condensed lab notebook in this README.

Command map

Run xlaxiom <command> --help for the complete interface. The documentation index 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 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

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 defines the required quality gates, reproducible-build expectations, and Trusted Publishing boundary. The roadmap distinguishes implemented capability from physical-hardware evidence that cannot be manufactured locally.