Support token dim in arbitrary location in dump comparator (#19455)

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

@@ -88,5 +88,114 @@ class TestExecuteAlignment:
         assert aligned.y.shape[1:] == (8,)
 
 
+class TestTokenDim:
+    """Tests for non-zero token_dim support."""
+
+    def _make_simple_plan(self, *, num_tokens: int) -> TokenAlignerPlan:
+        locator = TokenLocator(
+            steps=[0] * num_tokens,
+            token_index_in_step=list(range(num_tokens)),
+        )
+        return TokenAlignerPlan(locators=Pair(x=locator, y=locator))
+
+    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)
+        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)
+        for i in range(5):
+            assert torch.equal(
+                aligned.x.select(dim=1, index=i), tensor.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)
+        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)
+        for i in range(5):
+            assert torch.equal(
+                aligned.x.select(dim=2, index=i), tensor.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)
+        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)
+        for i in range(5):
+            assert torch.equal(aligned.x[i], tensor.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."""
+        torch.manual_seed(42)
+
+        plan = TokenAlignerPlan(
+            locators=Pair(
+                x=TokenLocator(steps=[], token_index_in_step=[]),
+                y=TokenLocator(steps=[], token_index_in_step=[]),
+            ),
+        )
+
+        # tensor shape [3, 5, 8], token_dim=1
+        tensors: dict[int, torch.Tensor] = {0: torch.randn(3, 5, 8)}
+        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]
+        assert aligned.x.shape == (3, 0, 8)
+        assert aligned.y.shape == (3, 0, 8)
+
+    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)
+        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)
+        for i in range(5):
+            assert torch.equal(
+                aligned.x.select(dim=2, index=i), tensor.select(dim=2, index=i)
+            )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))