Update token layout and cleanup printer in dump comparator (#19457)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py

@@ -24,7 +24,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
 from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
     normalize_parallel_info,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
@@ -61,7 +61,7 @@ def load_and_normalize_aux(
         if step_data:
             steps_data[step] = step_data
 
-    layout: str = plugin.detect_layout(steps_data)
+    layout: TokenLayout = plugin.detect_layout(steps_data)
 
     step_auxs: dict[int, TokenAlignerStepAux] = {
         step: plugin.compute_step_aux(step_data, layout=layout, step=step)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py

@@ -10,6 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     SGLangSeqId,
     TokenAlignerStepAux,
 )
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 
@@ -45,11 +46,11 @@ class _AuxFrameworkPlugin(ABC):
         return frozenset()
 
     @abstractmethod
-    def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: ...
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: ...
 
     @abstractmethod
     def compute_step_aux(
-        self, step_data: dict[str, object], *, layout: str, step: int
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
     ) -> TokenAlignerStepAux: ...
 
     @abstractmethod
@@ -89,11 +90,11 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
     def has_required_names(self, names: set[str]) -> bool:
         return "input_ids" in names and "seq_lens" in names
 
-    def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
-        return "thd"
+        return TokenLayout.T
 
     def compute_step_aux(
-        self, step_data: dict[str, object], *, layout: str, step: int
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
     ) -> TokenAlignerStepAux:
         input_ids = step_data["input_ids"]
         positions = step_data["positions"]
@@ -154,13 +155,13 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
     def has_required_names(self, names: set[str]) -> bool:
         return "input_ids" in names and "cu_seqlens_q" in names
 
-    def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
+    def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
         for step_data in raw.values():
             if (qkv_format := step_data.get("qkv_format")) is not None:
                 fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
                 if "bshd" in fmt.lower():
                     raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
-                return "thd"
+                return TokenLayout.T
 
             input_ids = step_data.get("input_ids")
             if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
@@ -175,10 +176,10 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
                 ),
             )
         )
-        return "thd"
+        return TokenLayout.T
 
     def compute_step_aux(
-        self, step_data: dict[str, object], *, layout: str, step: int
+        self, step_data: dict[str, object], *, layout: TokenLayout, step: int
     ) -> TokenAlignerStepAux:
         input_ids: torch.Tensor = step_data["input_ids"]
 

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py

@@ -5,6 +5,7 @@ from typing import NamedTuple, Union
 
 from pydantic import model_validator
 
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import (
     Pair,
     _check_equal_lengths,
@@ -50,7 +51,7 @@ class TokenAlignerGlobalAux:
 
     step_auxs: dict[int, TokenAlignerStepAux]
     framework: str  # "sglang" | "megatron"
-    layout: str  # "thd"
+    layout: TokenLayout
 
 
 class TokenLocator(_FrozenBase):
@@ -106,7 +107,7 @@ class TokenAlignerSeqsInfo(_FrozenBase):
     """All sequences for one side across all steps."""
 
     sequences: dict[SeqId, TokenAlignerSeqInfo]
-    layout: str
+    layout: TokenLayout
 
 
 class TokenAlignerPlan(_FrozenBase):

python/sglang/srt/debug_utils/comparator/dims.py

@@ -8,6 +8,11 @@ BATCH_DIM_NAME: str = "b"
 SEQ_DIM_NAME: str = "s"
 
 
+class TokenLayout(Enum):
+    T = "t"  # single flat token dim
+    BS = "bs"  # separate batch + seq dims, need collapse
+
+
 class ParallelAxis(Enum):
     TP = "tp"
     CP = "cp"

python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py

@@ -125,6 +125,7 @@ def _compute_diff(
         max_diff_coord=list(max_diff_coord),
         baseline_at_max=x_baseline[max_diff_coord].item(),
         target_at_max=x_target[max_diff_coord].item(),
+        diff_threshold=diff_threshold,
         passed=(
             rel_diff <= diff_threshold
             and max_abs_diff <= diff_threshold

python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py

@@ -69,13 +69,10 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list
 
 
 def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
-    marker = "✅" if diff.passed else "❌"
     return [
         prefix_text
-        + marker
-        + " "
         + "\t".join(
-            f"{name}={value}"
+            f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
             for name, value in [
                 ("rel_diff", diff.rel_diff),
                 ("max_abs_diff", diff.max_abs_diff),

python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py

@@ -1,86 +0,0 @@
-from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
-    DiffInfo,
-    TensorComparisonInfo,
-    TensorStats,
-)
-
-
-def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None:
-    baseline = info.baseline
-    target = info.target
-
-    dtype_marker = "" if baseline.dtype == target.dtype else "🟠"
-    print(
-        f"Raw "
-        f"[shape] {baseline.shape} vs {target.shape}\t"
-        f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}"
-    )
-
-    if info.unified_shape != baseline.shape:
-        print(
-            f"Unify shape: {baseline.shape} -> {info.unified_shape} "
-            f"(to match {target.shape})"
-        )
-
-    print(
-        f"After unify "
-        f"[shape] {info.unified_shape} vs {target.shape}\t"
-        f"[dtype] {baseline.dtype} vs {target.dtype}"
-    )
-
-    _print_stats_comparison(baseline=baseline.stats, target=target.stats)
-
-    if info.shape_mismatch:
-        print("⚠️ Shape mismatch")
-        return
-
-    if info.diff is not None:
-        _print_diff(
-            diff=info.diff,
-            diff_threshold=diff_threshold,
-        )
-
-    if info.diff_downcast is not None and info.downcast_dtype is not None:
-        _print_diff(
-            diff=info.diff_downcast,
-            diff_threshold=diff_threshold,
-            prefix_text=f"When downcast to {info.downcast_dtype}: ",
-        )
-
-    if baseline.sample is not None:
-        print(f"x_baseline(sample)={baseline.sample}")
-    if target.sample is not None:
-        print(f"x_target(sample)={target.sample}")
-
-
-def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None:
-    stat_names = list(TensorStats.model_fields.keys())
-    for stat_name in stat_names:
-        value_baseline = getattr(baseline, stat_name)
-        value_target = getattr(target, stat_name)
-        if value_baseline is None or value_target is None:
-            continue
-        print(
-            f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
-            f"(diff: {value_target - value_baseline:.4f})"
-        )
-
-
-def _print_diff(diff: DiffInfo, diff_threshold: float, prefix_text: str = "") -> None:
-    print(
-        prefix_text
-        + "\t".join(
-            f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
-            for name, value in [
-                ("rel_diff", diff.rel_diff),
-                ("max_abs_diff", diff.max_abs_diff),
-                ("mean_abs_diff", diff.mean_abs_diff),
-            ]
-        )
-    )
-
-    print(
-        f"max_abs_diff happens at coord={diff.max_diff_coord} with "
-        f"baseline={diff.baseline_at_max} "
-        f"target={diff.target_at_max}"
-    )

python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py

@@ -28,6 +28,7 @@ class DiffInfo(_StrictBase):
     max_diff_coord: list[int]
     baseline_at_max: float
     target_at_max: float
+    diff_threshold: float
     passed: bool
 
 

test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py

@@ -13,6 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     SGLangSeqId,
     TokenAlignerStepAux,
 )
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=15, suite="default", nightly=True)
@@ -34,7 +35,7 @@ class TestNormalizeSGLang:
         }
 
         result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
-            step_data, layout="thd", step=0
+            step_data, layout=TokenLayout.T, step=0
         )
 
         assert result.input_ids == [10, 20, 30]
@@ -51,7 +52,7 @@ class TestNormalizeSGLang:
         }
 
         result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
-            step_data, layout="thd", step=3
+            step_data, layout=TokenLayout.T, step=3
         )
         assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
 
@@ -65,7 +66,7 @@ class TestNormalizeSGLang:
         }
 
         result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
-            step_data, layout="thd", step=0
+            step_data, layout=TokenLayout.T, step=0
         )
         assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
 
@@ -81,7 +82,7 @@ class TestNormalizeMegatron:
         }
 
         result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
-            step_data, layout="thd", step=0
+            step_data, layout=TokenLayout.T, step=0
         )
 
         assert result.seq_lens == [3, 2]
@@ -94,7 +95,7 @@ class TestNormalizeMegatron:
         }
 
         result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
-            step_data, layout="thd", step=0
+            step_data, layout=TokenLayout.T, step=0
         )
 
         assert result.positions == [0, 1, 2, 0, 1]
@@ -108,7 +109,7 @@ class TestNormalizeMegatron:
         }
 
         result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
-            step_data, layout="thd", step=0
+            step_data, layout=TokenLayout.T, step=0
         )
 
         assert result.positions == [5, 6, 7, 8, 9]
@@ -121,7 +122,7 @@ class TestNormalizeMegatron:
         }
 
         result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
-            step_data, layout="thd", step=5
+            step_data, layout=TokenLayout.T, step=5
         )
         assert result.seq_ids == [
             PositionalSeqId(step=5, seq_index=0),

test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerStepAux,
     TokenLocator,
 )
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -52,7 +53,7 @@ class TestExecuteAlignment:
         side_aux = TokenAlignerGlobalAux(
             step_auxs={0: aux, 1: aux_step1},
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)

test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py

@@ -19,6 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerStepAux,
     TokenLocator,
 )
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -40,7 +41,7 @@ class TestBuildTokenIndexSGLangThd:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -75,7 +76,7 @@ class TestBuildTokenIndexSGLangThd:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -108,7 +109,7 @@ class TestBuildTokenIndexSGLangThd:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -132,7 +133,7 @@ class TestBuildTokenIndexSGLangThd:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -163,7 +164,7 @@ class TestBuildTokenIndexMegatronThd:
                 ),
             },
             framework="megatron",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -204,7 +205,7 @@ class TestBuildTokenIndexMegatronThd:
                 ),
             },
             framework="megatron",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index = build_seqs_info(side_aux)
@@ -389,7 +390,7 @@ class TestComputeAlignmentPlanCrossLayout:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
         side_aux_b = TokenAlignerGlobalAux(
             step_auxs={
@@ -401,7 +402,7 @@ class TestComputeAlignmentPlanCrossLayout:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index_a = build_seqs_info(side_aux_a)
@@ -428,7 +429,7 @@ class TestComputeAlignmentPlanCrossLayout:
                 ),
             },
             framework="sglang",
-            layout="thd",
+            layout=TokenLayout.T,
         )
         side_aux_b = TokenAlignerGlobalAux(
             step_auxs={
@@ -443,7 +444,7 @@ class TestComputeAlignmentPlanCrossLayout:
                 ),
             },
             framework="megatron",
-            layout="thd",
+            layout=TokenLayout.T,
         )
 
         index_a = build_seqs_info(side_aux_a)
@@ -467,7 +468,7 @@ def _int_to_seq_id(k: int) -> SeqId:
 def _make_index(
     *,
     sequences: dict[int, tuple[int, ...]],
-    layout: str = "thd",
+    layout: TokenLayout = TokenLayout.T,
 ) -> TokenAlignerSeqsInfo:
     """Create a TokenAlignerSeqsInfo from simplified input_ids-only specification."""
     records: dict[SeqId, TokenAlignerSeqInfo] = {}

test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py

@@ -35,6 +35,7 @@ def _make_diff(
     rel_diff: float = 0.0001,
     max_abs_diff: float = 0.0005,
     mean_abs_diff: float = 0.0002,
+    diff_threshold: float = 1e-3,
     passed: bool = True,
 ) -> DiffInfo:
     return DiffInfo(
@@ -44,6 +45,7 @@ def _make_diff(
         max_diff_coord=[2, 3],
         baseline_at_max=1.0,
         target_at_max=1.0005,
+        diff_threshold=diff_threshold,
         passed=passed,
     )
 
@@ -94,7 +96,7 @@ class TestFormatComparison:
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
+            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005"
         )
@@ -153,11 +155,11 @@ class TestFormatComparison:
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "❌ rel_diff=0.002\tmax_abs_diff=0.005\tmean_abs_diff=0.001\n"
+            "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005\n"
             "When downcast to torch.bfloat16: "
-            "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
+            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005"
         )
@@ -187,7 +189,7 @@ class TestFormatComparison:
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
+            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005"
         )
@@ -215,7 +217,7 @@ class TestFormatComparison:
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
+            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005\n"
             "x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
@@ -243,7 +245,7 @@ class TestFormatComparison:
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
+            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005"
         )

test/registered/debug_utils/comparator/tensor_comparator/test_printer.py

@@ -1,265 +0,0 @@
-import sys
-
-import pytest
-
-from sglang.srt.debug_utils.comparator.tensor_comparator.printer import (
-    print_comparison,
-)
-from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
-    DiffInfo,
-    TensorComparisonInfo,
-    TensorInfo,
-    TensorStats,
-)
-from sglang.test.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=10, suite="default", nightly=True)
-
-
-def _make_stats(
-    mean: float = 0.0,
-    std: float = 1.0,
-    min: float = -2.0,
-    max: float = 2.0,
-    p1: float | None = -1.8,
-    p5: float | None = -1.5,
-    p95: float | None = 1.5,
-    p99: float | None = 1.8,
-) -> TensorStats:
-    return TensorStats(
-        mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
-    )
-
-
-def _make_diff(
-    rel_diff: float = 0.0001,
-    max_abs_diff: float = 0.0005,
-    mean_abs_diff: float = 0.0002,
-    passed: bool = True,
-) -> DiffInfo:
-    return DiffInfo(
-        rel_diff=rel_diff,
-        max_abs_diff=max_abs_diff,
-        mean_abs_diff=mean_abs_diff,
-        max_diff_coord=[2, 3],
-        baseline_at_max=1.0,
-        target_at_max=1.0005,
-        passed=passed,
-    )
-
-
-def _make_tensor_info(
-    shape: list[int] | None = None,
-    dtype: str = "torch.float32",
-    stats: TensorStats | None = None,
-    sample: str | None = None,
-) -> TensorInfo:
-    return TensorInfo(
-        shape=shape if shape is not None else [4, 8],
-        dtype=dtype,
-        stats=stats if stats is not None else _make_stats(),
-        sample=sample,
-    )
-
-
-# Snapshot strings below are intentionally spelled out in full per test.
-# The shared skeleton (stats block, diff block) looks duplicated, but keeping
-# each test self-contained makes failures immediately readable without chasing
-# helper functions.  Do not extract common fragments.
-class TestPrintComparison:
-    def test_normal(self, capsys):
-        info = TensorComparisonInfo(
-            name="test",
-            baseline=_make_tensor_info(
-                stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0),
-            ),
-            target=_make_tensor_info(
-                stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001),
-            ),
-            unified_shape=[4, 8],
-            shape_mismatch=False,
-            diff=_make_diff(),
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "After unify [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n"
-            "[std] 1.0000 vs 1.0001 (diff: 0.0001)\n"
-            "[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
-            "[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
-            "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
-            "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
-            "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
-            "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-        )
-
-    def test_shape_mismatch(self, capsys):
-        info = TensorComparisonInfo(
-            name="mismatch",
-            baseline=_make_tensor_info(shape=[3, 4]),
-            target=_make_tensor_info(shape=[5, 6]),
-            unified_shape=[3, 4],
-            shape_mismatch=True,
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [3, 4] vs [5, 6]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "After unify [shape] [3, 4] vs [5, 6]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
-            "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
-            "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
-            "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
-            "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
-            "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
-            "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "⚠️ Shape mismatch\n"
-        )
-
-    def test_with_downcast(self, capsys):
-        info = TensorComparisonInfo(
-            name="downcast",
-            baseline=_make_tensor_info(),
-            target=_make_tensor_info(dtype="torch.bfloat16"),
-            unified_shape=[4, 8],
-            shape_mismatch=False,
-            diff=_make_diff(
-                rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False
-            ),
-            diff_downcast=_make_diff(
-                rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True
-            ),
-            downcast_dtype="torch.bfloat16",
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [4, 8] vs [4, 8]\t"
-            "[🟠dtype] torch.float32 vs torch.bfloat16\n"
-            "After unify [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.bfloat16\n"
-            "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
-            "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
-            "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
-            "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
-            "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
-            "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
-            "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-            "When downcast to torch.bfloat16: "
-            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-        )
-
-    def test_with_shape_unification(self, capsys):
-        info = TensorComparisonInfo(
-            name="unify",
-            baseline=_make_tensor_info(shape=[1, 1, 4, 8]),
-            target=_make_tensor_info(),
-            unified_shape=[4, 8],
-            shape_mismatch=False,
-            diff=_make_diff(),
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [1, 1, 4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "Unify shape: [1, 1, 4, 8] -> [4, 8] "
-            "(to match [4, 8])\n"
-            "After unify [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
-            "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
-            "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
-            "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
-            "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
-            "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
-            "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-        )
-
-    def test_with_samples(self, capsys):
-        info = TensorComparisonInfo(
-            name="samples",
-            baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"),
-            target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"),
-            unified_shape=[4, 8],
-            shape_mismatch=False,
-            diff=_make_diff(),
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "After unify [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
-            "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
-            "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
-            "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
-            "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
-            "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
-            "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-            "x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
-            "x_target(sample)=tensor([0.1, 0.3, ...])\n"
-        )
-
-    def test_none_quantiles(self, capsys):
-        stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
-
-        info = TensorComparisonInfo(
-            name="no_quantiles",
-            baseline=_make_tensor_info(stats=stats_no_quantiles),
-            target=_make_tensor_info(stats=stats_no_quantiles),
-            unified_shape=[4, 8],
-            shape_mismatch=False,
-            diff=_make_diff(),
-        )
-
-        print_comparison(info=info, diff_threshold=1e-3)
-
-        assert capsys.readouterr().out == (
-            "Raw [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "After unify [shape] [4, 8] vs [4, 8]\t"
-            "[dtype] torch.float32 vs torch.float32\n"
-            "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
-            "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
-            "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
-            "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
-            "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
-            "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005\n"
-        )
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/tensor_comparator/test_types.py

@@ -46,6 +46,7 @@ def _make_diff(**overrides) -> DiffInfo:
         max_diff_coord=[2, 3],
         baseline_at_max=1.0,
         target_at_max=1.0005,
+        diff_threshold=1e-3,
         passed=True,
     )
     defaults.update(overrides)
@@ -76,6 +77,7 @@ class TestStrictBase:
                 max_diff_coord=[0],
                 baseline_at_max=0.0,
                 target_at_max=0.0,
+                diff_threshold=1e-3,
                 passed=True,
                 extra_field=123,
             )

test/registered/debug_utils/comparator/test_model_validation.py

@@ -185,6 +185,7 @@ def _make_diff_info(*, passed: bool) -> DiffInfo:
         max_diff_coord=[0, 0],
         baseline_at_max=1.0,
         target_at_max=1.01,
+        diff_threshold=1e-3,
         passed=passed,
     )