Reorganize modules and pipeline in dump comparator (#19374)

python/sglang/srt/debug_utils/comparator/aligner/reorder.py

@@ -1,82 +0,0 @@
-from typing import Literal
-
-import torch
-
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
-from sglang.srt.debug_utils.comparator.utils import _FrozenBase
-
-
-class ZigzagToNaturalParams(_FrozenBase):
-    op: Literal["zigzag_to_natural"] = "zigzag_to_natural"
-    dim: int
-    cp_size: int
-
-
-ReorderParams = ZigzagToNaturalParams
-
-
-class ReorderPlan(_FrozenBase):
-    params: ReorderParams
-
-
-_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {"s"}
-
-
-def compute_reorder_plans(
-    dim_specs: list[DimSpec],
-    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
-) -> list[ReorderPlan]:
-    plans: list[ReorderPlan] = []
-
-    for dim_index, spec in enumerate(dim_specs):
-        if (
-            spec.ordering is not None
-            and spec.ordering != Ordering.NATURAL
-            and spec.parallel is not None
-        ):
-            if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
-                raise ValueError(
-                    f"Zigzag ordering is only supported on sequence dims "
-                    f"(bshd/sbhd format, dim name must be one of "
-                    f"{sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}), "
-                    f"but got dim name {spec.name!r} in {spec}"
-                )
-
-            assert spec.ordering == Ordering.ZIGZAG
-            axis_size: int = parallel_infos[0][spec.parallel].axis_size
-            plans.append(
-                ReorderPlan(
-                    params=ZigzagToNaturalParams(dim=dim_index, cp_size=axis_size),
-                )
-            )
-
-    return plans
-
-
-def execute_reorder_plan(
-    plan: ReorderPlan,
-    tensors: list[torch.Tensor],
-) -> list[torch.Tensor]:
-    return [
-        _reorder_zigzag_to_natural(
-            tensor, dim=plan.params.dim, cp_size=plan.params.cp_size
-        )
-        for tensor in tensors
-    ]
-
-
-def _reorder_zigzag_to_natural(
-    tensor: torch.Tensor, *, dim: int, cp_size: int
-) -> torch.Tensor:
-    """Undo CP zigzag interleaving, restoring natural chunk order.
-
-    Generalized from Megatron-LM _undo_attention_load_balancing
-    (megatron/core/ssm/mamba_context_parallel.py:360-373).
-    """
-    num_chunks: int = cp_size * 2
-    chunks: tuple[torch.Tensor, ...] = tensor.chunk(num_chunks, dim=dim)
-    order: list[int] = [2 * i for i in range(cp_size)] + [
-        num_chunks - 2 * i - 1 for i in range(cp_size)
-    ]
-    return torch.cat([chunks[i] for i in order], dim=dim)

python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py

@@ -0,0 +1,31 @@
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+
+
+def execute_reorderer_plan(
+    plan: ReordererPlan,
+    tensors: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    return [
+        _reorder_zigzag_to_natural(
+            tensor, dim=plan.params.dim, cp_size=plan.params.cp_size
+        )
+        for tensor in tensors
+    ]
+
+
+def _reorder_zigzag_to_natural(
+    tensor: torch.Tensor, *, dim: int, cp_size: int
+) -> torch.Tensor:
+    """Undo CP zigzag interleaving, restoring natural chunk order.
+
+    Generalized from Megatron-LM _undo_attention_load_balancing
+    (megatron/core/ssm/mamba_context_parallel.py:360-373).
+    """
+    num_chunks: int = cp_size * 2
+    chunks: tuple[torch.Tensor, ...] = tensor.chunk(num_chunks, dim=dim)
+    order: list[int] = [2 * i for i in range(cp_size)] + [
+        num_chunks - 2 * i - 1 for i in range(cp_size)
+    ]
+    return torch.cat([chunks[i] for i in order], dim=dim)

python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py

@@ -0,0 +1,39 @@
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
+    ReordererPlan,
+    ZigzagToNaturalParams,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
+
+_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {"s"}
+
+
+def compute_reorderer_plans(
+    dim_specs: list[DimSpec],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> list[ReordererPlan]:
+    plans: list[ReordererPlan] = []
+
+    for dim_index, spec in enumerate(dim_specs):
+        if (
+            spec.ordering is not None
+            and spec.ordering != Ordering.NATURAL
+            and spec.parallel is not None
+        ):
+            if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
+                raise ValueError(
+                    f"Zigzag ordering is only supported on sequence dims "
+                    f"(bshd/sbhd format, dim name must be one of "
+                    f"{sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}), "
+                    f"but got dim name {spec.name!r} in {spec}"
+                )
+
+            assert spec.ordering == Ordering.ZIGZAG
+            axis_size: int = parallel_infos[0][spec.parallel].axis_size
+            plans.append(
+                ReordererPlan(
+                    params=ZigzagToNaturalParams(dim=dim_index, cp_size=axis_size),
+                )
+            )
+
+    return plans

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

@@ -0,0 +1,16 @@
+from typing import Literal
+
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+
+class ZigzagToNaturalParams(_FrozenBase):
+    op: Literal["zigzag_to_natural"] = "zigzag_to_natural"
+    dim: int
+    cp_size: int
+
+
+ReordererParams = ZigzagToNaturalParams
+
+
+class ReordererPlan(_FrozenBase):
+    params: ReordererParams

python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py

@@ -1,56 +1,52 @@
 import torch
 
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     PickParams,
-    UnshardParams,
-    UnshardPlan,
+    UnsharderParams,
+    UnsharderPlan,
 )
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
-from sglang.srt.debug_utils.comparator.output_types import (
-    AnyWarning,
-    ReplicatedMismatchWarning,
-)
+from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 
 
-def execute_unshard_plan(
-    plan: UnshardPlan,
+def execute_unsharder_plan(
+    plan: UnsharderPlan,
     tensors: list[torch.Tensor],
-) -> tuple[list[torch.Tensor], list[AnyWarning]]:
-    all_warnings: list[AnyWarning] = []
+) -> list[torch.Tensor]:
     result: list[torch.Tensor] = []
 
     for group_idx, group in enumerate(plan.groups):
         group_tensors = [tensors[i] for i in group]
-        tensor, warnings = _apply_unshard(
+        tensor = _apply_unshard(
             plan.params,
             group_tensors,
             axis=plan.axis,
             group_index=group_idx,
         )
         result.append(tensor)
-        all_warnings.extend(warnings)
 
-    return result, all_warnings
+    return result
 
 
 def _apply_unshard(
-    params: UnshardParams,
+    params: UnsharderParams,
     ordered_tensors: list[torch.Tensor],
     *,
     axis: ParallelAxis,
     group_index: int,
-) -> tuple[torch.Tensor, list[AnyWarning]]:
+) -> torch.Tensor:
     if isinstance(params, PickParams):
-        warnings = _verify_replicated_group(
+        _verify_replicated_group(
             ordered_tensors,
             axis=axis,
             group_index=group_index,
         )
-        return ordered_tensors[0], warnings
+        return ordered_tensors[0]
 
     if isinstance(params, ConcatParams):
-        return torch.cat(ordered_tensors, dim=params.dim), []
+        return torch.cat(ordered_tensors, dim=params.dim)
 
     # Phase 2: ReduceSumParams, CpZigzagParams
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
@@ -61,14 +57,13 @@ def _verify_replicated_group(
     *,
     axis: ParallelAxis,
     group_index: int,
-) -> list[ReplicatedMismatchWarning]:
-    warnings: list[ReplicatedMismatchWarning] = []
+) -> None:
     baseline = ordered_tensors[0]
 
     for i in range(1, len(ordered_tensors)):
         other = ordered_tensors[i]
         if not torch.allclose(baseline, other, atol=1e-6):
-            warnings.append(
+            warning_sink.add(
                 ReplicatedMismatchWarning(
                     axis=axis.value,
                     group_index=group_index,
@@ -77,5 +72,3 @@ def _verify_replicated_group(
                     max_abs_diff=(baseline - other).abs().max().item(),
                 )
             )
-
-    return warnings

python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 
 _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")

python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py

@@ -1,12 +1,12 @@
 from collections import defaultdict
 from typing import NamedTuple
 
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     AxisInfo,
     ConcatParams,
     PickParams,
-    UnshardParams,
-    UnshardPlan,
+    UnsharderParams,
+    UnsharderPlan,
 )
 from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis
 
@@ -21,10 +21,10 @@ class _GroupResult(NamedTuple):
     projected_coords: _CoordsList
 
 
-def compute_unshard_plan(
+def compute_unsharder_plan(
     dim_specs: list[DimSpec],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
-) -> list[UnshardPlan]:
+) -> list[UnsharderPlan]:
     if not parallel_infos:
         raise ValueError("parallel_infos must not be empty")
 
@@ -51,20 +51,20 @@ def compute_unshard_plan(
         for info in parallel_infos
     ]
 
-    axis_and_params: list[tuple[ParallelAxis, UnshardParams]] = [
+    axis_and_params: list[tuple[ParallelAxis, UnsharderParams]] = [
         (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
     ] + [
         (axis, _resolve_unshard_params(spec=spec, dim_index=dim_index))
         for axis, (dim_index, spec) in sharded_axis_infos.items()
     ]
 
-    plans: list[UnshardPlan] = []
+    plans: list[UnsharderPlan] = []
     for axis, params in axis_and_params:
         result = _group_and_project(
             current_coords=current_coords,
             target_axis=axis,
         )
-        plans.append(UnshardPlan(axis=axis, params=params, groups=result.groups))
+        plans.append(UnsharderPlan(axis=axis, params=params, groups=result.groups))
         current_coords = result.projected_coords
 
     return plans
@@ -130,7 +130,7 @@ def _group_and_project(
     return _GroupResult(groups=groups, projected_coords=projected)
 
 
-def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnshardParams:
+def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnsharderParams:
     if spec.reduction is not None:
         raise NotImplementedError(
             f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"

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

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Annotated, Literal, Union
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 from sglang.srt.debug_utils.comparator.utils import _FrozenBase
@@ -12,6 +12,16 @@ class AxisInfo(_FrozenBase):
     axis_rank: int
     axis_size: int
 
+    @model_validator(mode="after")
+    def _validate_bounds(self) -> AxisInfo:
+        if self.axis_size <= 0:
+            raise ValueError(f"axis_size must be > 0, got {self.axis_size}")
+        if not (0 <= self.axis_rank < self.axis_size):
+            raise ValueError(
+                f"axis_rank must be in [0, {self.axis_size}), got {self.axis_rank}"
+            )
+        return self
+
 
 class ConcatParams(_FrozenBase):
     op: Literal["concat"] = "concat"
@@ -22,15 +32,15 @@ class PickParams(_FrozenBase):
     op: Literal["pick"] = "pick"
 
 
-UnshardParams = Annotated[
+UnsharderParams = Annotated[
     Union[ConcatParams, PickParams],
     Field(discriminator="op"),
 ]
 
 
-class UnshardPlan(_FrozenBase):
+class UnsharderPlan(_FrozenBase):
     axis: ParallelAxis
-    params: UnshardParams
+    params: UnsharderParams
     # groups[i] = indices in the input tensor list, which will be operated (e.g. concat) into i-th output tensor.
     #
     # Multistep example (CP=2, TP=2, 4 input tensors):

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

@@ -3,10 +3,10 @@ from typing import Annotated, Any, Literal, Union
 
 from pydantic import Discriminator, Field, TypeAdapter, model_validator
 
-from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
+from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
     format_comparison,
 )
-from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
     TensorComparisonInfo,
 )
 from sglang.srt.debug_utils.comparator.utils import _StrictBase

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

@@ -3,31 +3,36 @@ from typing import Any, Optional, Union
 
 import torch
 
-from sglang.srt.debug_utils.comparator.aligner.reorder import (
-    ReorderPlan,
-    compute_reorder_plans,
-    execute_reorder_plan,
+from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
+    execute_reorderer_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
-    execute_unshard_plan,
+from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
+    compute_reorderer_plans,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
+from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
+from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
+    execute_unsharder_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
     normalize_parallel_info,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
-    compute_unshard_plan,
+from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
+    compute_unsharder_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
 from sglang.srt.debug_utils.comparator.dims import parse_dims
 from sglang.srt.debug_utils.comparator.output_types import (
     AnyWarning,
     ComparisonRecord,
     SkipRecord,
 )
-from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    compare_tensor_pair,
+)
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 from sglang.srt.debug_utils.dump_loader import ValueWithMeta
 
-Plan = Union[UnshardPlan, ReorderPlan]
+Plan = Union[UnsharderPlan, ReordererPlan]
 
 
 def process_tensor_group(
@@ -38,6 +43,28 @@ def process_tensor_group(
     baseline_path: Path,
     target_path: Path,
     diff_threshold: float,
+) -> ComparisonRecord | SkipRecord:
+    with warning_sink.context() as collected_warnings:
+        return _process_tensor_group_raw(
+            name=name,
+            baseline_filenames=baseline_filenames,
+            target_filenames=target_filenames,
+            baseline_path=baseline_path,
+            target_path=target_path,
+            diff_threshold=diff_threshold,
+            collected_warnings=collected_warnings,
+        )
+
+
+def _process_tensor_group_raw(
+    *,
+    name: str,
+    baseline_filenames: list[str],
+    target_filenames: list[str],
+    baseline_path: Path,
+    target_path: Path,
+    diff_threshold: float,
+    collected_warnings: list[AnyWarning],
 ) -> ComparisonRecord | SkipRecord:
     b_tensors = _load_tensors(baseline_filenames, baseline_path)
     t_tensors = _load_tensors(target_filenames, target_path)
@@ -51,22 +78,21 @@ def process_tensor_group(
     t_extracted = _extract_tensors(t_tensors)
     del b_tensors, t_tensors
 
-    b_tensor, b_warns = _execute_plans(b_extracted, b_plans)
-    t_tensor, t_warns = _execute_plans(t_extracted, t_plans)
-    all_warnings: list[AnyWarning] = b_warns + t_warns
+    b_tensor = _execute_plans(b_extracted, b_plans)
+    t_tensor = _execute_plans(t_extracted, t_plans)
 
     if b_tensor is None or t_tensor is None:
         reason = "baseline_load_failed" if b_tensor is None else "target_load_failed"
-        return SkipRecord(name=name, reason=reason, warnings=all_warnings)
+        return SkipRecord(name=name, reason=reason, warnings=collected_warnings)
 
-    info = compare_tensors(
+    info = compare_tensor_pair(
         x_baseline=b_tensor,
         x_target=t_tensor,
         name=name,
         diff_threshold=diff_threshold,
     )
 
-    return ComparisonRecord(**info.model_dump(), warnings=all_warnings)
+    return ComparisonRecord(**info.model_dump(), warnings=collected_warnings)
 
 
 def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
@@ -96,13 +122,13 @@ def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]:
     dim_specs = parse_dims(dims_str)
     parallel_infos = [normalize_parallel_info(meta) for meta in metas]
 
-    unshard_plans = compute_unshard_plan(
+    unsharder_plans = compute_unsharder_plan(
         dim_specs=dim_specs, parallel_infos=parallel_infos
     )
-    reorder_plans = compute_reorder_plans(
+    reorderer_plans = compute_reorderer_plans(
         dim_specs=dim_specs, parallel_infos=parallel_infos
     )
-    return [*unshard_plans, *reorder_plans]
+    return [*unsharder_plans, *reorderer_plans]
 
 
 def _extract_tensors(
@@ -114,32 +140,30 @@ def _extract_tensors(
 def _execute_plans(
     tensors: list[torch.Tensor],
     plans: list[Plan],
-) -> tuple[Optional[torch.Tensor], list[AnyWarning]]:
+) -> Optional[torch.Tensor]:
     if not tensors:
-        return None, []
+        return None
 
     if not plans:
         if len(tensors) != 1:
-            return None, []
-        return tensors[0], []
+            return None
+        return tensors[0]
 
-    warnings: list[AnyWarning] = []
     current = tensors
     for plan in plans:
-        current, new_warnings = _execute_plan(current, plan)
-        warnings.extend(new_warnings)
+        current = _execute_plan(current, plan)
 
     assert len(current) == 1
-    return current[0], warnings
+    return current[0]
 
 
 def _execute_plan(
     tensors: list[torch.Tensor],
     plan: Plan,
-) -> tuple[list[torch.Tensor], list[AnyWarning]]:
-    if isinstance(plan, UnshardPlan):
-        return execute_unshard_plan(plan, tensors)
-    elif isinstance(plan, ReorderPlan):
-        return execute_reorder_plan(plan, tensors), []
+) -> list[torch.Tensor]:
+    if isinstance(plan, UnsharderPlan):
+        return execute_unsharder_plan(plan, tensors)
+    elif isinstance(plan, ReordererPlan):
+        return execute_reorderer_plan(plan, tensors)
     else:
         raise NotImplementedError(f"Unknown {plan=}")

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

@@ -0,0 +1,3 @@
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    compare_tensor_pair,
+)

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

@@ -2,13 +2,14 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
     DiffInfo,
     TensorComparisonInfo,
     TensorInfo,
     TensorStats,
 )
 from sglang.srt.debug_utils.comparator.utils import (
+    Pair,
     argmax_coord,
     calc_rel_diff,
     compute_smaller_dtype,
@@ -20,7 +21,7 @@ QUANTILE_NUMEL_THRESHOLD = 10_000_000
 SAMPLE_DIFF_THRESHOLD = 1e-3
 
 
-def compare_tensors(
+def compare_tensor_pair(
     x_baseline: torch.Tensor,
     x_target: torch.Tensor,
     name: str = "",
@@ -66,7 +67,7 @@ def compare_tensors(
 
         if baseline_original_dtype != target_original_dtype:
             downcast_dtype = compute_smaller_dtype(
-                baseline_original_dtype, target_original_dtype
+                Pair(x=baseline_original_dtype, y=target_original_dtype)
             )
             if downcast_dtype is not None:
                 diff_downcast = _compute_diff(

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

@@ -1,4 +1,4 @@
-from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
     DiffInfo,
     TensorComparisonInfo,
     TensorStats,

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

@@ -1,4 +1,4 @@
-from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
+from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
     DiffInfo,
     TensorComparisonInfo,
     TensorStats,

python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py

@@ -1 +0,0 @@
-from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors

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

@@ -40,13 +40,13 @@ def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]:
 
 
 def compute_smaller_dtype(
-    dtype_a: torch.dtype, dtype_b: torch.dtype
+    dtypes: Pair[torch.dtype],
 ) -> Optional[torch.dtype]:
     info_dict = {
         (torch.float32, torch.bfloat16): torch.bfloat16,
         # ... add more ...
     }
-    return info_dict.get((dtype_a, dtype_b)) or info_dict.get((dtype_b, dtype_a))
+    return info_dict.get((dtypes.x, dtypes.y)) or info_dict.get((dtypes.y, dtypes.x))
 
 
 def try_unify_shape(x: torch.Tensor, target_shape: torch.Size) -> torch.Tensor: