scripts/triposplat/validate_gaussian_decoder_onnx.py

#!/usr/bin/env python3
"""Validate Gaussian decoder ONNX features against official TripoSplat PyTorch.

The command feeds identical float32 ``points`` and ``cond`` tensors to upstream
``decoder.gs`` and Python ONNX Runtime, reports detailed errors, and exits nonzero
unless every raw feature value satisfies the configured allclose gate.  Gaussian
offset construction, representation scaling, biases, and activations are deliberately
outside this validation boundary.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from decoder_onnx_common import (
    COND_SHAPE,
    FEATURES_SHAPE,
    OFFICIAL_REPOSITORY_URL,
    POINTS_SHAPE,
    array_summary,
    choose_torch_device,
    comparison_metrics,
    default_tolerances,
    deterministic_gaussian_inputs,
    load_gaussian_fixture,
    load_graph_contract,
    run_pytorch_reference,
    validate_array,
    validate_gaussian_inputs,
    validate_with_ort,
    validation_runtime_report,
    write_json,
)


COMPONENT = "gaussian_decoder"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--triposplat-repo",
        type=Path,
        required=True,
        help=f"Local clone of {OFFICIAL_REPOSITORY_URL}.",
    )
    parser.add_argument(
        "--weights",
        type=Path,
        required=True,
        help="Official decoder checkpoint used for the ONNX export.",
    )
    parser.add_argument(
        "--onnx",
        type=Path,
        required=True,
        help="Exported fixed-shape Gaussian decoder ONNX graph.",
    )
    parser.add_argument(
        "--fixture",
        type=Path,
        help="Optional NPZ with points, cond and optionally features.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=20260714,
        help="Synthetic input seed when --fixture is omitted (default: %(default)s).",
    )
    parser.add_argument(
        "--precision",
        choices=("auto", "fp16", "fp32"),
        default="auto",
        help="Expected internal graph precision; auto reads metadata (default: %(default)s).",
    )
    parser.add_argument(
        "--device",
        choices=("cpu", "mps", "cuda", "auto"),
        default="cpu",
        help="Device for the official PyTorch reference (default: %(default)s).",
    )
    parser.add_argument(
        "--provider",
        action="append",
        dest="providers",
        help=(
            "Python ONNX Runtime provider. Repeat for fallback order "
            "(default: CPUExecutionProvider)."
        ),
    )
    parser.add_argument("--atol", type=float, help="Absolute allclose tolerance.")
    parser.add_argument("--rtol", type=float, help="Relative allclose tolerance.")
    parser.add_argument("--report", type=Path, help="Optional JSON report destination.")
    parser.add_argument(
        "--save-fixture",
        type=Path,
        help="Optional NPZ destination containing inputs and official features.",
    )
    parser.add_argument(
        "--session-threads",
        type=int,
        default=0,
        help="ORT intra-op threads; zero leaves the runtime default (default: %(default)s).",
    )
    args = parser.parse_args()
    if args.atol is not None and args.atol < 0:
        parser.error("--atol must be non-negative")
    if args.rtol is not None and args.rtol < 0:
        parser.error("--rtol must be non-negative")
    if args.session_threads < 0:
        parser.error("--session-threads must be non-negative")
    return args


def validate(args: argparse.Namespace) -> dict[str, Any]:
    try:
        import numpy as np
        import onnx
        import onnxruntime as ort
        import torch
    except ImportError as exc:
        raise SystemExit(
            "Missing validation dependency. Install a PyTorch-supported Python "
            "version and run `python -m pip install -r "
            "scripts/triposplat/requirements.txt`. "
            f"Original error: {exc}"
        ) from exc

    graph_path = args.onnx.expanduser().resolve()
    metadata, artifacts, precision = load_graph_contract(
        onnx,
        graph_path,
        COMPONENT,
        args.precision,
    )
    default_atol, default_rtol = default_tolerances(precision)
    atol = args.atol if args.atol is not None else default_atol
    rtol = args.rtol if args.rtol is not None else default_rtol

    if args.fixture:
        points, cond, fixture_expected = load_gaussian_fixture(args.fixture)
        input_source = str(args.fixture.expanduser().resolve())
    else:
        points, cond = deterministic_gaussian_inputs(args.seed)
        fixture_expected = None
        input_source = f"deterministic synthetic inputs (seed={args.seed})"
    validate_gaussian_inputs(points, cond)

    device = choose_torch_device(torch, args.device)
    print(f"Running official PyTorch decoder.gs ({precision} internal) on {device}")
    reference, torch_duration_ms = run_pytorch_reference(
        torch=torch,
        component=COMPONENT,
        triposplat_repo=args.triposplat_repo,
        weights=args.weights,
        device=device,
        internal_precision=precision,
        inputs=(points, cond),
    )
    validate_array(reference, "PyTorch features", FEATURES_SHAPE)

    providers = args.providers or ["CPUExecutionProvider"]
    print(f"Running Python ONNX Runtime with providers {providers}")
    candidate, ort_duration_ms, session = validate_with_ort(
        ort=ort,
        graph_path=graph_path,
        providers=providers,
        threads=args.session_threads,
        output_name="features",
        feeds={"points": points, "cond": cond},
    )
    validate_array(candidate, "ONNX Runtime features", FEATURES_SHAPE)

    metrics = comparison_metrics(reference, candidate, atol=atol, rtol=rtol)
    fixture_metrics = None
    passed = bool(metrics["passed"])
    if fixture_expected is not None:
        fixture_metrics = comparison_metrics(
            fixture_expected,
            reference,
            atol=atol,
            rtol=rtol,
        )
        passed = bool(passed and fixture_metrics["passed"])

    report: dict[str, Any] = {
        "passed": passed,
        "component": COMPONENT,
        "contract": {
            "points": {
                "shape": list(POINTS_SHAPE),
                "dtype": "float32",
                "range": [0.0, 1.0],
            },
            "cond": {"shape": list(COND_SHAPE), "dtype": "float32"},
            "features": {"shape": list(FEATURES_SHAPE), "dtype": "float32"},
            "internal_precision": precision,
            "gaussians_per_point": 32,
        },
        "graph_metadata": metadata,
        "artifacts": [
            {"path": str(path), "bytes": path.stat().st_size} for path in artifacts
        ],
        "inputs": {
            "source": input_source,
            "points": array_summary(points),
            "cond": array_summary(cond),
        },
        "comparison": metrics,
        "fixture_reference_comparison": fixture_metrics,
        "outputs": {
            "pytorch": array_summary(reference),
            "onnxruntime": array_summary(candidate),
        },
        "tolerance": {"atol": atol, "rtol": rtol},
        "measured_single_run_ms": {
            "pytorch": torch_duration_ms,
            "onnxruntime_python": ort_duration_ms,
            "note": "Measured validation runs only; these are not browser WebGPU benchmarks.",
        },
        "runtime": validation_runtime_report(
            torch=torch,
            ort=ort,
            device=device,
            session=session,
            graph_path=graph_path,
        ),
    }

    if args.save_fixture:
        fixture_path = args.save_fixture.expanduser().resolve()
        fixture_path.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(
            fixture_path,
            points=np.ascontiguousarray(points, dtype=np.float32),
            cond=np.ascontiguousarray(cond, dtype=np.float32),
            features=np.ascontiguousarray(reference, dtype=np.float32),
            metadata=np.asarray(
                json.dumps(
                    {
                        "component": COMPONENT,
                        "internal_precision": precision,
                        "source": "official TripoSplat decoder.gs PyTorch reference",
                    },
                    sort_keys=True,
                )
            ),
        )
        report["saved_fixture"] = str(fixture_path)
        print(f"Wrote browser parity fixture {fixture_path}")
    if args.report:
        write_json(args.report, report)
        print(f"Wrote report {args.report.expanduser().resolve()}")

    print(json.dumps(report["comparison"], indent=2, sort_keys=True))
    print(
        "PASS: official PyTorch and Python ONNX Runtime features agree"
        if report["passed"]
        else "FAIL: Gaussian features are outside the requested tolerances"
    )
    return report


def main() -> None:
    report = validate(parse_args())
    if not report["passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()