ExportPass opens a second FakeTensorMode when every input is a lifted constant, aborting lowering

[msluszniak]

🐛 Describe the bug

ExportPass retraces under a brand new FakeTensorMode when every placeholder in the graph is a lifted constant, leaving the resulting graph holding fake tensors from two different modes. detect_fake_mode() rejects that downstream, and because the assertion escapes the pass manager it aborts the entire lowering.

Repro (no backend, no model download):

import torch
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
from executorch.exir.pass_base import ExportPass

mode = FakeTensorMode(allow_non_fake_inputs=True)

def const_fake(value):
    """A FakeTensor carrying its real value, as constant propagation makes."""
    real = torch.tensor(value)
    with mode:
        meta = torch.empty(real.shape, dtype=real.dtype, device="meta")
    return FakeTensor(mode, meta, torch.device("cpu"), constant=real)

graph = torch.fx.Graph()
lhs = graph.placeholder("c_lhs")
rhs = graph.placeholder("c_rhs")
lhs.meta["val"] = const_fake(2.0)
rhs.meta["val"] = const_fake(3.0)
mul = graph.call_function(torch.ops.aten.mul.Tensor, (lhs, rhs))
with mode:
    mul.meta["val"] = torch.ops.aten.mul.Tensor(lhs.meta["val"], rhs.meta["val"])
out = graph.output((mul,))
out.meta["val"] = (mul.meta["val"],)
gm = torch.fx.GraphModule(torch.nn.Module(), graph)

result = ExportPass()(gm)
modes = {
    id(t.fake_mode)
    for n in result.graph_module.graph.nodes
    for t in ((v := n.meta.get("val")) if isinstance(v, (list, tuple)) else [v])
    if isinstance(t, FakeTensor)
}
print(len(modes))   # 2 on main, expected 1

Root cause

_ExportPassBase.call picks the mode by scanning self.inputs() for a FakeTensor:

fake_tensor_mode = None
for i in inputs:
    if isinstance(i, FakeTensor):
        ...
        fake_tensor_mode = i.fake_mode
if fake_tensor_mode is None:
    fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True)

but inputs() unwraps a constant-carrying fake tensor to its real tensor first:

def extract_input(node):
    if "val" in node.meta:
        fake = node.meta["val"]
        if hasattr(fake, "constant") and fake.constant is not None:
            return fake.constant
        return fake

So when every placeholder is such a constant, the scan sees no FakeTensor, a fresh mode is opened, and the retraced nodes end up in it while the untouched placeholders keep the old one.

How it shows up in practice

Partitioners produce exactly this shape of submodule: an input-independent subgraph whose scalar operands become lifted constant placeholders when it is split out. A sine positional embedding is enough. Lowering RF-DETR nano to the Vulkan backend fails with:

Exception: An error occurred when running the 'FuseBatchNormPass' pass after the following passes: []

which names an unrelated pass and says nothing about fake modes. The delegate submodule involved:

mode A: _lifted_tensor_constant4, _lifted_tensor_constant5, _lifted_tensor_constant6   (placeholders)
mode B: aten_arange_start_step, aten_div_tensor_mode, aten_mul_tensor, ..., output

Reduced end-to-end repro against VulkanPartitioner:

class M(torch.nn.Module):
    def forward(self, x):
        t = torch.arange(128, dtype=torch.float32)
        t = 10000.0 ** (2 * torch.div(t, 2, rounding_mode="floor") / 128)
        return x + t

Expected behavior

The pass should retrace under the graph's existing fake mode rather than inventing a second one.

PR: #22310

Versions

ExecuTorch main @ c27baa8 (also reproduces on the v1.4.1 branch). Python 3.10, torch 2.13.0, macOS 15.5 / arm64.