Use named tensors in dump comparator (#19458)

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

@@ -28,14 +28,18 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=15, suite="default", nightly=True)
 
 
+def _named(tensor: torch.Tensor, names: list[str]) -> torch.Tensor:
+    return tensor.refine_names(*names)
+
+
 class TestExecuteAlignment:
     """Tests for token alignment execution."""
 
     def test_thd_vs_thd_identity(self):
         """Two identical thd sides produce element-wise equal aligned tensors."""
         torch.manual_seed(42)
-        hidden_step0 = torch.randn(5, 8)  # 5 tokens, hidden_dim=8
-        hidden_step1 = torch.randn(2, 8)  # 2 tokens
+        hidden_step0 = torch.randn(5, 8).refine_names("t", "h")
+        hidden_step1 = torch.randn(2, 8).refine_names("t", "h")
 
         aux = TokenAlignerStepAux(
             input_ids=[10, 20, 30, 40, 50],
@@ -78,7 +82,7 @@ class TestExecuteAlignment:
             ),
         )
 
-        tensors = {0: torch.randn(5, 8)}
+        tensors = {0: torch.randn(5, 8).refine_names("t", "h")}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors)
         )
@@ -102,58 +106,58 @@ class TestTokenDim:
     def test_token_dim_nonzero(self) -> None:
         """tensor shape [3, 5, 8], token_dim=1 -> token dim stays at dim 1."""
         torch.manual_seed(42)
-        tensor: torch.Tensor = torch.randn(3, 5, 8)
+        tensor: torch.Tensor = _named(torch.randn(3, 5, 8), ["a", "t", "h"])
         plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5)
 
         tensors: dict[int, torch.Tensor] = {0: tensor}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan,
             tensor_of_step_pair=Pair(x=tensors, y=tensors),
-            token_dims=Pair(x=1, y=1),
         )
 
         assert aligned.x.shape == (3, 5, 8)
         assert torch.equal(aligned.x, aligned.y)
+        plain: torch.Tensor = tensor.rename(None)
         for i in range(5):
             assert torch.equal(
-                aligned.x.select(dim=1, index=i), tensor.select(dim=1, index=i)
+                aligned.x.select(dim=1, index=i), plain.select(dim=1, index=i)
             )
 
     def test_token_dim_last(self) -> None:
         """tensor shape [3, 8, 5], token_dim=2 -> token dim stays at dim 2."""
         torch.manual_seed(42)
-        tensor: torch.Tensor = torch.randn(3, 8, 5)
+        tensor: torch.Tensor = _named(torch.randn(3, 8, 5), ["a", "h", "t"])
         plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5)
 
         tensors: dict[int, torch.Tensor] = {0: tensor}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan,
             tensor_of_step_pair=Pair(x=tensors, y=tensors),
-            token_dims=Pair(x=2, y=2),
         )
 
         assert aligned.x.shape == (3, 8, 5)
+        plain: torch.Tensor = tensor.rename(None)
         for i in range(5):
             assert torch.equal(
-                aligned.x.select(dim=2, index=i), tensor.select(dim=2, index=i)
+                aligned.x.select(dim=2, index=i), plain.select(dim=2, index=i)
             )
 
     def test_token_dim_zero(self) -> None:
         """token_dim=0 selects along first dimension (standard t-h-d layout)."""
         torch.manual_seed(42)
-        tensor: torch.Tensor = torch.randn(5, 8)
+        tensor: torch.Tensor = _named(torch.randn(5, 8), ["t", "h"])
         plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5)
 
         tensors: dict[int, torch.Tensor] = {0: tensor}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan,
             tensor_of_step_pair=Pair(x=tensors, y=tensors),
-            token_dims=Pair(x=0, y=0),
         )
 
         assert aligned.x.shape == (5, 8)
+        plain: torch.Tensor = tensor.rename(None)
         for i in range(5):
-            assert torch.equal(aligned.x[i], tensor.select(dim=0, index=i))
+            assert torch.equal(aligned.x[i], plain.select(dim=0, index=i))
 
     def test_zero_matched_tokens_nonzero_token_dim(self) -> None:
         """Empty plan with token_dim=1 produces correct empty shape."""
@@ -166,12 +170,12 @@ class TestTokenDim:
             ),
         )
 
-        # tensor shape [3, 5, 8], token_dim=1
-        tensors: dict[int, torch.Tensor] = {0: torch.randn(3, 5, 8)}
+        tensors: dict[int, torch.Tensor] = {
+            0: _named(torch.randn(3, 5, 8), ["a", "t", "h"])
+        }
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan,
             tensor_of_step_pair=Pair(x=tensors, y=tensors),
-            token_dims=Pair(x=1, y=1),
         )
 
         # token dim (dim 1) set to 0, other dims preserved -> [3, 0, 8]
@@ -181,20 +185,22 @@ class TestTokenDim:
     def test_high_rank_tensor(self) -> None:
         """tensor shape [2, 3, 5, 4, 8] (a b t c d), token_dim=2 -> stays at dim 2."""
         torch.manual_seed(42)
-        tensor: torch.Tensor = torch.randn(2, 3, 5, 4, 8)
+        tensor: torch.Tensor = _named(
+            torch.randn(2, 3, 5, 4, 8), ["a", "x", "t", "c", "h"]
+        )
         plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5)
 
         tensors: dict[int, torch.Tensor] = {0: tensor}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan,
             tensor_of_step_pair=Pair(x=tensors, y=tensors),
-            token_dims=Pair(x=2, y=2),
         )
 
         assert aligned.x.shape == (2, 3, 5, 4, 8)
+        plain: torch.Tensor = tensor.rename(None)
         for i in range(5):
             assert torch.equal(
-                aligned.x.select(dim=2, index=i), tensor.select(dim=2, index=i)
+                aligned.x.select(dim=2, index=i), plain.select(dim=2, index=i)
             )