From 4274b174e92e87b058f7fe102c75ae5d6820f005 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 00:56:08 +0800
Subject: [PATCH] feat(tools): add ORT vs PyTorch parity check for exports

---
 tools/export_onnx.py | 89 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/tools/export_onnx.py b/tools/export_onnx.py
index 054979d..c665381 100644
--- a/tools/export_onnx.py
+++ b/tools/export_onnx.py
@@ -169,3 +169,92 @@ def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path:
     )
     logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3)
     return out_path
+
+
+def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None:
+    """Verify ONNX flow model matches PyTorch output on random input."""
+    import onnxruntime as ort
+    from tools.models import TrajectoryFlowModel
+
+    cfg = json.loads((ckpt_dir / "train_config.json").read_text())
+    seq_len = int(cfg["seq_len"])
+    cond_dim = int(cfg.get("cond_dim", 3))
+
+    model = TrajectoryFlowModel(
+        seq_len=seq_len,
+        d_model=int(cfg["d_model"]),
+        nhead=int(cfg["nhead"]),
+        num_layers=int(cfg["num_layers"]),
+        dim_feedforward=int(cfg["dim_feedforward"]),
+        cond_dim=cond_dim,
+        dropout=0.0,
+    )
+    model.load_state_dict(
+        torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
+    )
+    model.eval()
+
+    torch.manual_seed(42)
+    np.random.seed(42)
+    x = torch.randn(2, seq_len, 3, dtype=torch.float32)
+    t = torch.tensor([0.0, 0.5], dtype=torch.float32)
+    cond = torch.randn(2, cond_dim, dtype=torch.float32)
+
+    with torch.no_grad():
+        torch_out = model(x, t, cond).numpy()
+
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    ort_out = sess.run(
+        ["v"],
+        {
+            "x_t": x.numpy(),
+            "t": t.numpy(),
+            "cond": cond.numpy(),
+        },
+    )[0]
+
+    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
+        max_diff = float(np.abs(torch_out - ort_out).max())
+        raise RuntimeError(
+            f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
+        )
+    logger.info("Flow model parity OK (atol=%.0e)", _ATOL)
+
+
+def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None:
+    """Verify ONNX scroll decoder matches PyTorch decoder output."""
+    import onnxruntime as ort
+    from tools.scroll.models import ScrollCVAE
+
+    cfg = json.loads((ckpt_dir / "scroll_config.json").read_text())
+    seq_len = int(cfg["seq_len"])
+    latent_dim = int(cfg["latent_dim"])
+    cond_dim = int(cfg["cond_dim"])
+
+    full = ScrollCVAE(
+        seq_len=seq_len,
+        latent_dim=latent_dim,
+        hidden=int(cfg["hidden"]),
+        cond_dim=cond_dim,
+    )
+    full.load_state_dict(
+        torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
+    )
+    full.eval()
+
+    torch.manual_seed(7)
+    z = torch.randn(2, latent_dim, dtype=torch.float32)
+    cond = torch.randn(2, cond_dim, dtype=torch.float32)
+
+    with torch.no_grad():
+        torch_out = full.decode(z, cond).numpy()
+
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0]
+
+    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
+        max_diff = float(np.abs(torch_out - ort_out).max())
+        raise RuntimeError(
+            f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
+        )
+    logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL)