complex_graph_rewrite: AttributeError 'float' object has no attribute 'name' when graph contains scalar literals near view_as_complex

[robbysun]

Bug description

torch_tensorrt.dynamo.compile crashes with AttributeError: 'float' object has no attribute 'name' when the exported graph contains both view_as_complex/view_as_real ops and a scalar constant (e.g. from x * 1.0 / math.sqrt(dim)) nearby.

The complex_graph_detection lowering pass's match_complex_mul filter iterates over match.nodes_map.values(), assuming every value is an fx.Node. However, torch.fx.subgraph_rewriter can include Python scalar literals (floats) in nodes_map, causing the crash.

Minimal reproducer

import math
import torch
import torch.nn as nn
import torch_tensorrt

class ComplexRoPEWithScale(nn.Module):
    def __init__(self, dim: int, seq_len: int):
        super().__init__()
        freqs = torch.polar(
            torch.ones(seq_len, dim // 2),
            torch.arange(seq_len * dim // 2, dtype=torch.float).reshape(seq_len, dim // 2),
        )
        self.register_buffer("freqs_cis", freqs)
        self.scale = 1.0 / math.sqrt(dim)  # becomes a float literal in the fx graph

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # view_as_complex -> mul -> view_as_real  (triggers complex_graph_detection)
        xc = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
        out = torch.view_as_real(xc * self.freqs_cis).flatten(-2)
        # scalar multiply — introduces a Python float constant near the complex subgraph
        return (out * self.scale).to(x.dtype)

B, L, D = 2, 16, 32
x = torch.randn(B, L, D, device="cuda", dtype=torch.bfloat16)
model = ComplexRoPEWithScale(dim=D, seq_len=L).cuda().bfloat16().eval()

with torch.no_grad():
    ep = torch.export.export(model, (x,))

trt_ep = torch_tensorrt.dynamo.compile(
    ep,
    inputs=[x],
    enabled_precisions={torch.bfloat16},
    use_explicit_typing=False,
    device=torch_tensorrt.Device(gpu_id=0),
)

Full traceback

Traceback (most recent call last):
  File "repro.py", line 34, in <module>
  File ".../torch_tensorrt/dynamo/_compiler.py", line 764, in compile
    gm = post_lowering(gm, settings)
  File ".../torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py", line 137, in post_lowering
    gm = ATEN_POST_LOWERING_PASSES(gm, settings)
  File ".../torch_tensorrt/dynamo/lowering/passes/pass_manager.py", line 135, in __call__
    out = _pass(out, settings)
  File ".../torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py", line 360, in complex_graph_detection
    complex_graph_rewriter.rewrite_subgraph_nodes(complex_subgraphs)
  File ".../torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py", line 221, in rewrite_subgraph_nodes
    nodes = torch.fx.subgraph_rewriter.replace_pattern_with_filters(...)
  File ".../torch/fx/subgraph_rewriter.py", line 310, in <genexpr>
    match_filter(m, original_graph, pattern_graph)
  File ".../torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py", line 217, in match_complex_mul
    if original_node.name == node.name:
       ^^^^^^^^^^^^^^^^^^
AttributeError: 'float' object has no attribute 'name'

Root cause

In complex_graph_rewrite.py, rewrite_subgraph_nodes (line ~216) defines match_complex_mul as:

def match_complex_mul(match, original_graph, pattern_graph) -> bool:
    for original_node in match.nodes_map.values():
        if original_node.name == node.name:   # ← crashes if value is a float
            return True
    return False

torch.fx.subgraph_rewriter.replace_pattern_with_filters passes ignore_literals=True, but match.nodes_map can still contain Python scalar values (floats, ints) that were captured as constants in the pattern. Accessing .name on a Python float raises AttributeError.

Fix

Guard the attribute access:

def match_complex_mul(match, original_graph, pattern_graph) -> bool:
    for original_node in match.nodes_map.values():
        if isinstance(original_node, torch.fx.Node) and original_node.name == node.name:
            return True
    return False

Environment

torch-tensorrt	2.10.0
torch	2.10.0+cu128
tensorrt	10.14.1.48.post1
CUDA	13.0
GPU	NVIDIA GeForce RTX 4080
Python	3.12

Context

Encountered while compiling a SAM3 ViT image encoder (rotary position encoding uses view_as_complex; the attention scaling factor 1/sqrt(head_dim) introduces the float literal). The _patch_vit_rope_real workaround (replacing complex RoPE with equivalent real arithmetic) avoids this path entirely, but the crash should be fixed in the library.

Discovered and diagnosed by @robbysun with assistance from Claude Code.

[narendasan]

@robbysun are you still seeing this issue with more recent versions of Torch-TRT?

[robbysun]

@robbysun are you still seeing this issue with more recent versions of Torch-TRT?

Confirmed on v2.10.0 + PyTorch 2.10.0. I haven't yet tested v2.12.x — upgrading requires a full PyTorch version bump which I haven't done. The complex_graph_rewrite pass appears to have landed in v2.12.0, so that would be the version to test against. I'll update the issue once I've had a chance to test it.