From 6ed39774901fa64c43020827de00841668b67a48 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 10:47:58 -0700
Subject: [PATCH 01/36] fix(converters): clamp out-of-range static slice bounds
 like Python

In Python and PyTorch, a slice bound past either end of a dimension is
clamped to that end. On a dimension of size 10, `x[-15:]` is the whole
dimension and `x[:-15]` is empty.

When the sliced dimension has a static size, the slice converter
normalized its bounds with `get_positive_dim`, which wraps a negative
value with a modulo. That is right for a dimension index but not for a
slice bound: `x[-15:]` became `x[5:]` and `x[:-15]` became `x[:5]`. The
engine built and ran, and silently returned the wrong rows. A slice
whose start lands after its stop (an empty slice in PyTorch) could fail
the engine build with a negative output length.

This shows up in real models. In transformers, the sliding-window
layers of the dynamic key/value cache (used by Gemma2 and Gemma3, for
example) store `keys[:, :, -window + 1:]`. With a static prompt shorter
than the window, the stored cache lost the first prompt tokens for most
prompt lengths, in every sliding-window layer.

Clamp static bounds the same way the converter already clamps a start
computed at runtime, and floor the output length at zero. Literal
bounds on a dynamic dimension take a different branch and are not
changed here.
---
 py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py | 9 ++++++---
 tests/py/dynamo/conversion/test_slice_aten.py         | 3 +++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
index a89113b6c3f..86a90ebd1ed 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
@@ -157,8 +157,11 @@ def slice_op(  # TODO: This should be slice not whatever is in base
     output_shape = list(input.shape)
 
     if input.shape[dim] != -1 and isinstance(start, int) and isinstance(stop, int):
-        start = get_positive_dim(start, input.shape[dim])
-        stop = get_positive_dim(stop, input.shape[dim])
+        # Clamp out-of-range bounds like Python slicing instead of wrapping them
+        # like a dimension index: x[-15:] on a size-10 dim is the whole dim.
+        dim_size = input.shape[dim]
+        start = max(start + dim_size, 0) if start < 0 else min(start, dim_size)
+        stop = max(stop + dim_size, 0) if stop < 0 else min(stop, dim_size)
         start_slice[dim] = start
     else:
         # the start and stop or None is dynamic along dim or or start or stop is an ITensor
@@ -272,7 +275,7 @@ def slice_op(  # TODO: This should be slice not whatever is in base
             layer.set_input(3, stride_slice_tensor)
             return layer.get_output(0)
 
-    output_shape[dim] = math.ceil((stop - start) / step)
+    output_shape[dim] = max(math.ceil((stop - start) / step), 0)
     return slice(
         ctx, target, source_ir, name, input, start_slice, output_shape, stride_slice
     )
diff --git a/tests/py/dynamo/conversion/test_slice_aten.py b/tests/py/dynamo/conversion/test_slice_aten.py
index bdf171f5fbb..7a8f2beb82c 100644
--- a/tests/py/dynamo/conversion/test_slice_aten.py
+++ b/tests/py/dynamo/conversion/test_slice_aten.py
@@ -19,6 +19,9 @@ class TestSliceConverter(DispatchTestCase):
             ("slice_dim_start_stop_step_max_int", 2, 0, 2**63 - 1, 1),
             ("slice_dim_start_stop_step_past_end", 2, 0, 2048, 1),
             ("slice_dim_start_stop_step_none", 2, None, None, 1),
+            ("slice_dim_start_before_begin", 1, -15, None, 1),
+            ("slice_dim_stop_before_begin", 1, 0, -15, 1),
+            ("slice_dim_start_after_stop", 1, 7, 3, 1),
         ]
     )
     def test_slice(self, _, dim, start, stop, step):

From 201ae83faa0fa85267f3a41abfb6495df56ae5b2 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 10:50:57 -0700
Subject: [PATCH 02/36] fix(partitioning): avoid quadratic frontier scans in
 the fast partitioner

The fast (adjacency) partitioner can take hours on large exported graphs.
On a vision-language-action model whose lowered graph has 345k nodes, it
did not finish in 75 minutes.

TRTPartitioner.starter_nodes puts every callable node in the frontier.
The inherited _SplitterBase.put_nodes_into_subgraphs then looks through
that frontier for a node whose dependencies are all placed, once per
placed node, so the work can grow with the square of the node count.

Override put_nodes_into_subgraphs with a dependency counter: each node
keeps the number of its dependencies not yet placed, and becomes ready
when the count reaches zero. The time is now proportional to the nodes
plus their dependency entries. A subgraph still takes every node of its
kind that becomes ready, so the subgraphs hold the same nodes as before.
Only the order inside a subgraph can change, and tagging and splitting do
not depend on it.

On that model, building the subgraphs now takes about 1 second; on just
its first 80k nodes the old loop took about 5 minutes. On a synthetic
62k-node graph it goes from 134 s to 0.1 s, with the same subgraphs and
the same split module.

The new test checks that the subgraphs match the base class result and
that the dependency sets are read a bounded number of times.
---
 .../partitioning/_adjacency_partitioner.py    | 63 +++++++++++++++++++
 .../test_000_fast_partitioning.py             | 48 ++++++++++++++
 2 files changed, 111 insertions(+)

diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
index 0b46c420286..4d9ae72b0dc 100644
--- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
+++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import logging
+from collections import defaultdict, deque
 from typing import Collection, Dict, List, Optional, Set, Tuple
 
 import torch
@@ -10,6 +11,7 @@
 from torch.fx.passes.splitter_base import (
     FxNetAccFusionsFinder,
     FxNetAccNodesFinder,
+    FxNetSplitterInternalError,
     Subgraph,
     _SplitterBase,
     _SplitterSettingBase,
@@ -412,6 +414,67 @@ def starter_nodes(self) -> Tuple[NodeSet, NodeSet]:
 
         return starter_non_acc_nodes, starter_acc_nodes
 
+    def put_nodes_into_subgraphs(self) -> List[Subgraph]:
+        """Builds the same subgraphs as _SplitterBase.put_nodes_into_subgraphs
+
+        starter_nodes puts every callable node in the frontier, and the base
+        class can rescan that whole frontier for a ready node at every step,
+        which is quadratic on large graphs. Counting the unplaced dependencies
+        of each node finds ready nodes in time proportional to the nodes plus
+        their dependency entries. Each subgraph still takes every node of its
+        kind whose dependencies are placed, so the node sets match. Fusion
+        groups need no extra step: every callable node is counted from the
+        start, and their constraints are already in deps.
+        """
+        pending = {
+            node: len(self.deps[node])
+            for node in self.module.graph.nodes
+            if node.op in CALLABLE_NODE_OPS
+        }
+        dependents: Dict[torch.fx.Node, List[torch.fx.Node]] = defaultdict(list)
+        for node in pending:
+            for dep in self.deps[node]:
+                dependents[dep].append(node)
+
+        ready = {
+            is_acc: deque(
+                node
+                for node, count in pending.items()
+                if count == 0 and (node in self.acc_nodes) == is_acc
+            )
+            for is_acc in (False, True)
+        }
+        acc_subgraph = not ready[False]
+        current_subgraph_nodes: List[torch.fx.Node] = []
+        subgraphs: List[Subgraph] = []
+        remaining = len(pending)
+        while remaining:
+            if not ready[acc_subgraph]:
+                if not current_subgraph_nodes:
+                    raise FxNetSplitterInternalError("Subgraph can't be empty")
+                subgraphs.append(
+                    Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes)
+                )
+                acc_subgraph = not acc_subgraph
+                current_subgraph_nodes = []
+                continue
+
+            node = ready[acc_subgraph].popleft()
+            current_subgraph_nodes.append(node)
+            remaining -= 1
+            for user in dependents[node]:
+                pending[user] -= 1
+                if pending[user] == 0:
+                    ready[user in self.acc_nodes].append(user)
+
+        if current_subgraph_nodes:
+            subgraphs.append(
+                Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes)
+            )
+        if not subgraphs:
+            raise FxNetSplitterInternalError("Couldn't create subgraphs")
+        return subgraphs
+
 
 def partition(
     gm: torch.fx.GraphModule,
diff --git a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py
index 6a666ca3dba..76dc3e8e337 100644
--- a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py
+++ b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py
@@ -1,18 +1,33 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from collections import defaultdict
 from copy import deepcopy
 
 import numpy as np
 import torch
 import torch_tensorrt
 from parameterized import parameterized
+from torch.fx.passes.splitter_base import _SplitterBase
+from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo import partitioning
+from torch_tensorrt.dynamo.partitioning._adjacency_partitioner import (
+    OpSupportTester,
+    TRTPartitioner,
+)
 
 from ..testing_utilities import lower_graph_testing
 
 
+class _CountingDeps(defaultdict):
+    reads = 0
+
+    def __getitem__(self, node):
+        self.reads += 1
+        return super().__getitem__(node)
+
+
 class TestFastPartitioning(TestCase):
     def test_partition_fully_supported_one_op(self):
         class FullySupportedOneOp(torch.nn.Module):
@@ -183,6 +198,39 @@ def forward(self, x, y):
             "Certain operators are set to run in Torch, expected 1 segment",
         )
 
+    def test_put_nodes_into_subgraphs_matches_base_without_rescans(self):
+        class Interleaved(torch.nn.Module):
+            def forward(self, x, y):
+                for i in range(60):
+                    a = torch.ops.aten.add.Tensor(x, y)
+                    b = torch.ops.aten.mul.Tensor(a, y)
+                    if i % 6 == 0:
+                        b = torch.ops.aten.abs.default(b)
+                    x = torch.ops.aten.sub.Tensor(b, a)
+                return x
+
+        partitioner = TRTPartitioner(
+            torch.fx.symbolic_trace(Interleaved()),
+            OpSupportTester(torch_executed_ops={"torch.ops.aten.abs.default"}),
+            min_block_size=1,
+        )
+        partitioner.deps = _CountingDeps(set, partitioner.deps)
+        subgraphs = partitioner.put_nodes_into_subgraphs()
+        reads = partitioner.deps.reads
+
+        expected = _SplitterBase.put_nodes_into_subgraphs(partitioner)
+        self.assertEqual(
+            [(s.is_acc, set(s.nodes)) for s in subgraphs],
+            [(s.is_acc, set(s.nodes)) for s in expected],
+        )
+        self.assertEqual(len(subgraphs), 21)
+
+        # The base loop can reread the dependencies of all waiting nodes at every step.
+        num_nodes = len(
+            [n for n in partitioner.module.graph.nodes if n.op in CALLABLE_NODE_OPS]
+        )
+        self.assertLessEqual(reads, 4 * num_nodes)
+
     @parameterized.expand(
         [
             (["torch.nn.modules.conv.Conv2d"], 2),

From e87304e1e5c4cf5fc5f095b8a4b29f199618bb04 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 12:24:22 -0700
Subject: [PATCH 03/36] fix(lowering): keep the attention output layout when
 attention is decomposed

With decompose_attention=True, scaled_dot_product_attention is replaced by
matmul and softmax, which return a contiguous tensor. The fused kernels lay
their output out differently: flash and cuDNN attention follow the query's
dimension order, and memory-efficient attention returns a transpose of an
(N, L, H, E) buffer.

An exported graph can depend on that layout. In the common pattern

    out = F.scaled_dot_product_attention(q, k, v)
    out = out.transpose(1, 2).contiguous().view(n, l, -1)

the transpose of the fused output is often already contiguous, so export
does not record the .contiguous(). After the decomposition the transpose
is not contiguous any more, and torch_tensorrt.dynamo.compile fails in
run_decompositions with "Cannot view a tensor with shape ... and strides
... as ...". The Florence-2 vision tower hits this.

The decomposition cannot tell which backend produced the layout, but the
attention node records it. A new pre-lowering pass, which runs only when
attention is decomposed, copies each attention output whose recorded
dimension order is not the default into that order (permute, contiguous
clone, permute back, then a clone that keeps this layout). It also visits
subgraphs: torch.cond branches and torch.no_grad() or autocast regions.
The original attention node becomes that last clone, so graph outputs,
in-place ops on the output and in-place updates through its views keep
working on the same tensor.

torch.compile with decompose_attention=True still fails the same way: its
backend does not run the pre-lowering passes, so it needs a separate fix.
---
 .../lowering/passes/_aten_lowering_pass.py    |  2 +
 .../preserve_attention_output_layout.py       | 89 +++++++++++++++++++
 .../lowering/test_aten_lowering_passes.py     | 31 +++++++
 3 files changed, 122 insertions(+)
 create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
index e58c0fef06d..1bcde1ec1f1 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
@@ -36,6 +36,7 @@
 from .lower_associative_scan import lower_associative_scan
 from .normalize_negative_slice_stop import normalize_negative_slice_stop  # noqa: F401
 from .pass_manager import DynamoPassManager
+from .preserve_attention_output_layout import preserve_attention_output_layout
 from .remove_alias import remove_alias
 from .remove_assert_nodes import remove_assert_nodes
 from .remove_detach import remove_detach
@@ -51,6 +52,7 @@
     remove_detach,
     remove_assert_nodes,
     rule_based_autocast,
+    preserve_attention_output_layout,
 ]
 
 
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py
new file mode 100644
index 00000000000..11060ff3167
--- /dev/null
+++ b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py
@@ -0,0 +1,89 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import logging
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensor
+from torch_tensorrt.dynamo._settings import CompilationSettings
+
+logger = logging.getLogger(__name__)
+
+
+def _preserve_layouts(graph: torch.fx.Graph) -> int:
+    count = 0
+    for node in list(graph.nodes):
+        if node.target != torch.ops.aten.scaled_dot_product_attention.default:
+            continue
+        val = node.meta.get("val")
+        # Under torch.vmap the recorded value is a batched tensor, not a FakeTensor,
+        # and those nodes lower without this pass.
+        if not isinstance(val, FakeTensor) or val.is_contiguous():
+            continue
+        dim_order = val.dim_order()
+        if dim_order == tuple(range(val.dim())):
+            continue
+
+        inverse = [dim_order.index(dim) for dim in range(val.dim())]
+        with graph.inserting_before(node):
+            attention = graph.call_function(node.target, node.args, node.kwargs)
+            permuted = graph.call_function(
+                torch.ops.aten.permute.default, (attention, list(dim_order))
+            )
+            copied = graph.call_function(
+                torch.ops.aten.clone.default,
+                (permuted,),
+                {"memory_format": torch.contiguous_format},
+            )
+            restored = graph.call_function(
+                torch.ops.aten.permute.default, (copied, inverse)
+            )
+        for new_node in (attention, permuted, copied, restored):
+            new_node.meta.update(
+                {key: value for key, value in node.meta.items() if key != "tensor_meta"}
+            )
+        # Recorded as the decomposition returns it, so a second run leaves it alone.
+        attention.meta["val"] = val.contiguous()
+        permuted.meta["val"] = attention.meta["val"].permute(dim_order)
+        copied.meta["val"] = permuted.meta["val"].contiguous()
+        restored.meta["val"] = copied.meta["val"].permute(inverse)
+        # The original node becomes a copy that keeps the recorded layout, so every
+        # user, graph outputs and in-place updates included, keeps reading the same
+        # node. A copy rather than a view, so in-place ops such as detach_() work.
+        node.target = torch.ops.aten.clone.default
+        node.args = (restored,)
+        node.kwargs = {}
+        count += 1
+    return count
+
+
+def preserve_attention_output_layout(
+    gm: torch.fx.GraphModule, settings: CompilationSettings
+) -> torch.fx.GraphModule:
+    """Keep the exported memory layout of attention outputs that will be decomposed.
+
+    The fused attention kernels return their output in a layout that depends on the
+    backend and the query, and export drops a ``.contiguous()`` that is a no-op on it.
+    The attention decomposition returns a contiguous tensor instead, so a later
+    ``view`` fails when the graph is retraced. Copying each output into the layout
+    recorded at export keeps those views valid.
+    """
+    if not settings.decompose_attention:
+        return gm
+
+    count = 0
+    # Submodules too: attention in torch.cond branches and in torch.no_grad() or
+    # autocast regions lives in subgraphs, nested ones included.
+    for module in gm.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            continue
+        changed = _preserve_layouts(module.graph)
+        if changed:
+            # No dead code elimination: the graph is not functionalized yet, and it
+            # would drop nodes with no users such as the vmap nesting calls.
+            module.recompile()
+            count += changed
+
+    if count:
+        logger.debug("Preserved the output layout of %d attention nodes", count)
+    return gm
diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
index 80d531f29b4..9448224d9d8 100644
--- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py
+++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
@@ -454,6 +454,37 @@ def test_eliminates_noop_sym_min_int64_max(self):
         self.assertEqual(output_node.args[0], (x, x))
 
 
+class TestPreserveAttentionOutputLayout(TestCase):
+    def test_view_after_decomposed_attention(self):
+        # Fused attention returns this transposed query's layout, so export drops the
+        # .contiguous() below and records a view that only works on that layout. The
+        # attention output must stay a tensor of its own (detach_), the in-place
+        # update through the view must reach it, and the no_grad region puts all of
+        # it in a subgraph.
+        class ViewAfterAttention(torch.nn.Module):
+            def forward(self, query, key, value):
+                with torch.no_grad():
+                    query, key, value = (t.transpose(1, 2) for t in (query, key, value))
+                    out = torch.nn.functional.scaled_dot_product_attention(
+                        query, key, value
+                    )
+                    out.detach_()
+                    flat = out.transpose(1, 2).contiguous()
+                    flat = flat.view(flat.shape[0], flat.shape[1], -1)
+                    flat.add_(1)
+                return out, flat
+
+        inputs = [torch.randn(2, 16, 4, 32).half().cuda() for _ in range(3)]
+        model = ViewAfterAttention().cuda()
+        ep = torch.export.export(model, tuple(inputs))
+        trt_module = torch_tensorrt.dynamo.compile(
+            ep, inputs, min_block_size=1, decompose_attention=True
+        )
+        torch.testing.assert_close(
+            model(*inputs), trt_module(*inputs), rtol=1e-2, atol=1e-2
+        )
+
+
 class TestRewriteEfficientAttention(TestCase):
     def test_force_causal_efficient_attention(self):
         class RewriteEfficientAttention(torch.nn.Module):

From 576f3332da85bb06f12e183bf551aadd13dfb0b3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 14:22:12 -0700
Subject: [PATCH 04/36] fix(lowering): keep the torch.cond branches that
 post_lowering returns

post_lowering lowers each torch.cond branch by calling itself on the
branch module, but it ignored the module that call returned. Most passes
edit the graph in place, so this went unnoticed. The complex
decomposition does not: it retraces the branch and returns a new module.
That module was dropped and the parent kept the original branch, complex
ops included, which TensorRT cannot convert. A conditional with complex
math inside a branch therefore stayed in PyTorch, and
require_full_compilation=True failed.

post_lowering now installs the returned module in place of the branch,
under every name the parent uses for it, and lowers a branch shared by
several conditionals once. The cond node was traced against the old
branch, so the new module is installed only if its inputs and outputs
keep their dtypes and memory layout: strides, storage offset and storage
size. A branch that returns a complex tensor changes its outputs to the
real layout, so it is left as it was. So is a branch whose retraced
output comes back in a different memory layout, for example contiguous
where the original was transposed, because code after the conditional
can read the output through its strides. Symbolic values that cannot be
proven equal count as different, so data-dependent shapes keep the
original branch instead of failing. A branch that takes a complex
operand, such as a complex buffer, still runs in PyTorch as before.
---
 .../lowering/passes/_aten_lowering_pass.py    | 51 +++++++++-
 tests/py/dynamo/models/test_cond.py           | 99 +++++++++++++++++++
 2 files changed, 147 insertions(+), 3 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
index 1bcde1ec1f1..a9733825f3a 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
@@ -3,9 +3,11 @@
 
 import logging
 import operator
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import Any, Callable, List, Optional, Sequence, Union
 
 import torch
+import torch.utils._pytree as pytree
+from torch.fx.experimental.symbolic_shapes import statically_known_true
 from torch_tensorrt._utils import is_tegra_platform
 from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo.lowering.passes._FakeTensorUpdater import FakeTensorUpdater
@@ -195,12 +197,55 @@ def post_lowering(
     if fake_mode is not None:
         fake_tensor_updater.incremental_update(fake_mode)
 
-    for sub in iter_cond_subgraphs(gm):
-        post_lowering(sub, settings)
+    # A branch shared by several cond nodes is lowered once.
+    for branch in dict.fromkeys(iter_cond_subgraphs(gm)):
+        lowered = post_lowering(branch, settings)
+        # A pass can return a new module instead of editing the branch in place.
+        # The cond node was traced against the old one, so take the new module only
+        # if its inputs and outputs keep their dtypes and memory layout.
+        if lowered is branch or not _same_io(lowered, branch):
+            continue
+        for name, module in list(gm.named_modules(remove_duplicate=False)):
+            if module is branch:
+                gm.set_submodule(name, lowered)
 
     return gm
 
 
+def _io_values(gm: torch.fx.GraphModule) -> List[Any]:
+    """The meta value of each graph input and output, in order."""
+
+    def value(node: Any) -> Any:
+        return node.meta.get("val") if isinstance(node, torch.fx.Node) else None
+
+    placeholders = [value(n) for n in gm.graph.nodes if n.op == "placeholder"]
+    outputs = [value(n) for n in pytree.tree_leaves(gm.graph.output_node().args)]
+    return placeholders + outputs
+
+
+def _same_io(a: torch.fx.GraphModule, b: torch.fx.GraphModule) -> bool:
+    """Whether two modules take and return tensors of the same dtype and layout."""
+
+    def same(x: Any, y: Any) -> bool:
+        if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor):
+            return not isinstance(x, torch.Tensor) and not isinstance(y, torch.Tensor)
+        # A symbolic value that cannot be proven equal counts as different.
+        return (
+            x.dtype == y.dtype
+            and x.dim() == y.dim()
+            and statically_known_true(
+                x.untyped_storage().nbytes() == y.untyped_storage().nbytes()
+            )
+            and statically_known_true(x.storage_offset() == y.storage_offset())
+            and all(
+                statically_known_true(p == q) for p, q in zip(x.stride(), y.stride())
+            )
+        )
+
+    a_values, b_values = _io_values(a), _io_values(b)
+    return len(a_values) == len(b_values) and all(map(same, a_values, b_values))
+
+
 def pre_export_lowering(
     ep: torch.export.ExportedProgram,
     settings: CompilationSettings = CompilationSettings(),
diff --git a/tests/py/dynamo/models/test_cond.py b/tests/py/dynamo/models/test_cond.py
index ee5ce8c5dfb..aeedb561be3 100644
--- a/tests/py/dynamo/models/test_cond.py
+++ b/tests/py/dynamo/models/test_cond.py
@@ -57,6 +57,71 @@ def forward(self, x, predicate):
         )
 
 
+class _ComplexCond(nn.Module):
+    def forward(self, x, predicate):
+        def rotate(value):
+            z = torch.view_as_complex(value.reshape(2, 2, 2))
+            return torch.view_as_real(z * z).reshape(2, 4)
+
+        return torch.cond(predicate, rotate, lambda value: value - 1, (x,))
+
+
+class _TransposedComplexCond(nn.Module):
+    def forward(self, x, predicate):
+        def rotate(value):
+            z = torch.view_as_complex(value)
+            return torch.view_as_real((z * z).transpose(0, 1))
+
+        result = torch.cond(
+            predicate, rotate, lambda value: (value - 1).transpose(0, 1), (x,)
+        )
+        return torch.as_strided(result, (24,), (1,))
+
+
+class _SlicedComplexCond(nn.Module):
+    def __init__(self, rows):
+        super().__init__()
+        self.rows = rows
+
+    def forward(self, x, predicate):
+        def rotate(value):
+            z = torch.view_as_complex(value)
+            return torch.view_as_real((z * z)[self.rows])
+
+        result = torch.cond(
+            predicate, rotate, lambda value: (value - 1)[self.rows], (x,)
+        )
+        return torch.as_strided(result, (24,), (1,), 0)
+
+
+class _DataDependentComplexCond(nn.Module):
+    def forward(self, x, n, predicate):
+        k = n.item()
+        torch._check(k >= 1)
+        torch._check(k <= x.shape[0])
+
+        def rotate(value):
+            z = torch.view_as_complex(value)
+            return torch.view_as_real((z * z).transpose(0, 1))
+
+        return torch.cond(
+            predicate,
+            rotate,
+            lambda value: (value - 1).transpose(0, 1),
+            (x[:k].contiguous(),),
+        )
+
+
+class _ComplexOutputCond(nn.Module):
+    def forward(self, x, predicate):
+        return torch.cond(
+            predicate,
+            lambda value: torch.view_as_complex(value + 1),
+            lambda value: torch.view_as_complex(value - 1),
+            (x,),
+        )
+
+
 def _export_lowered(mod: nn.Module, inputs: tuple) -> torch.fx.GraphModule:
     settings = CompilationSettings()
     exported = torch.export.export(mod.eval(), inputs)
@@ -73,6 +138,40 @@ def test_cond_is_supported_when_branches_are_convertible(self):
         gm = _export_lowered(_AddSubCond(), (x, pred))
         self.assertTrue(_cond_node(gm) in CONVERTERS)
 
+    def test_cond_keeps_branches_returned_by_lowering(self):
+        # The complex lowering returns a new branch module rather than editing it.
+        x = torch.randn(2, 4)
+        pred = torch.tensor(True)
+        gm = _export_lowered(_ComplexCond(), (x, pred))
+        self.assertTrue(_cond_node(gm) in CONVERTERS)
+
+    def test_cond_keeps_branch_when_lowering_changes_output_layout(self):
+        pred = torch.tensor(True)
+        cases = [
+            # The lowered branch is contiguous, so as_strided reads another order.
+            (_TransposedComplexCond(), (torch.randn(3, 4, 2), pred)),
+            # The lowered branches own smaller storages, one also at another offset.
+            (_SlicedComplexCond(slice(1, None)), (torch.randn(4, 3, 2), pred)),
+            (_SlicedComplexCond(slice(None, -1)), (torch.randn(4, 3, 2), pred)),
+            # Data-dependent strides cannot be proven equal.
+            (
+                _DataDependentComplexCond(),
+                (torch.randn(8, 4, 2), torch.tensor(5), pred),
+            ),
+        ]
+        for model, inputs in cases:
+            with self.subTest(type(model).__name__):
+                gm = _export_lowered(model, inputs)
+                torch.testing.assert_close(gm(*inputs), model(*inputs))
+
+    def test_cond_keeps_branch_when_lowering_changes_output_dtype(self):
+        # The lowered branch returns the real [..., 2] layout instead of complex.
+        x = torch.randn(2, 2, 2)
+        pred = torch.tensor(True)
+        model = _ComplexOutputCond()
+        gm = _export_lowered(model, (x, pred))
+        torch.testing.assert_close(gm(x, pred), model(x, pred))
+
     def test_cond_falls_back_when_branch_has_unsupported_op(self):
         x = torch.randn(1, 4).abs() + 0.1
         pred = torch.tensor(True)

From 111c3c529f8d59eac42228d336b545e8f04ea570 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 12:53:46 -0700
Subject: [PATCH 05/36] fix(lowering): let the complex decomposition retrace
 offloaded constants

With offload_module_to_cpu=True, constant folding keeps the folded
constants on the CPU to save GPU memory, while their metadata still
records the GPU. complex_decomposition_adapter then retraces the graph
with those CPU tensors next to GPU activations, and the retrace fails:

  FakeTensorDeviceMismatchError: Expected all tensors to be on the same
  device, but found at least two devices, cuda:0 and cpu!

The adapter falls back to the legacy complex pass. That pass fails the
same way on a small rotary embedding model, and on a large model it did
not finish within 45 minutes. So models with complex rotary embeddings
compiled with offload_module_to_cpu either fail or take the slow path.

The adapter now copies each offloaded constant to the device its
metadata records while the retrace runs, then puts the CPU tensors back,
both in the module it was given and in the module it returns, and frees
the copies. Where constants live after lowering does not change.

A successful retrace returns a new module, which no longer holds the
parameters that only constant folding read. compile() and
cross_compile_for_windows() now offload the module lowering started from
as well as the lowered one, so those parameters still leave the GPU.
---
 py/torch_tensorrt/dynamo/_compiler.py         |  8 +++
 .../passes/complex_decomposition_adapter.py   | 68 +++++++++++++++++--
 .../dynamo/lowering/test_constant_folding.py  | 48 +++++++++++++
 3 files changed, 118 insertions(+), 6 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index a32301aa910..3bbcf33eed3 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -428,12 +428,16 @@ def cross_compile_for_windows(
 
     # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires
     # module parameters to still be on GPU, so we must not deallocate before this call.
+    unlowered_gm = gm
     gm = post_lowering(gm, settings)
     logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
+        # Complex lowering can return a new module without the weights that only
+        # constant folding read, so offload the module lowering started from too.
+        deallocate_module(unlowered_gm)
         deallocate_module(gm)
         logger.info(
             "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False"
@@ -909,6 +913,7 @@ def compile(
 
     # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires
     # module parameters to still be on GPU, so we must not deallocate before this call.
+    unlowered_gm = gm
     gm = post_lowering(gm, settings)
     logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
     logger.debug("Lowered Input graph: " + str(gm.graph))
@@ -921,6 +926,9 @@ def compile(
 
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
+        # Complex lowering can return a new module without the weights that only
+        # constant folding read, so offload the module lowering started from too.
+        deallocate_module(unlowered_gm)
         deallocate_module(gm)
         logger.info(
             "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False"
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
index e35265461e0..aab92106ef0 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
@@ -27,8 +27,10 @@
 such as ``aten.log2``/``aten.log10`` as of torch 2.14.dev).
 """
 
+import gc
 import logging
-from typing import Any
+from contextlib import contextmanager
+from typing import Any, Iterator
 
 import torch
 from torch._export.utils import _detect_fake_mode_from_gm
@@ -116,11 +118,12 @@ def complex_decomposition_adapter(
             if any(free_symbols(arg) for arg in flat_args)
             else decompose_complex_in_graph
         )
-        decomposed_gm = decompose(
-            gm,
-            flat_args,
-            decompositions=_trt_decomposition_table(settings, gm),
-        )
+        with _offloaded_constants_on_device(gm) as offloaded:
+            decomposed_gm = decompose(
+                gm,
+                flat_args,
+                decompositions=_trt_decomposition_table(settings, gm),
+            )
     except Exception as e:
         # decompose_complex_in_graph is upstream, experimental PyTorch code
         # (torch._functorch._aot_autograd.complex_decomposition) with
@@ -138,6 +141,11 @@ def complex_decomposition_adapter(
         return complex_graph_detection(gm, settings)
     gm.graph._codegen = original_codegen  # the caller's module, as it was
     gm.recompile()
+    _restore_offloaded_constants(decomposed_gm, offloaded)
+    if offloaded:
+        # The retrace leaves the device copies in reference cycles. Free them now
+        # rather than holding them through the engine build.
+        gc.collect()
     gm = decomposed_gm
 
     # The retrace brings back alias nodes. aten.alias has no converter, so the
@@ -239,6 +247,54 @@ def _graph_has_complex(gm: GraphModule) -> bool:
     return False
 
 
+@contextmanager
+def _offloaded_constants_on_device(
+    gm: GraphModule,
+) -> Iterator[dict[int, torch.Tensor]]:
+    """Hand the retrace each CPU-offloaded constant on the device it was computed on.
+
+    With offload_module_to_cpu, constant folding keeps the folded values on the CPU
+    while their metadata still records the device that computed them. The retrace
+    runs the graph, so a CPU constant next to device activations fails with a device
+    mismatch. Copy each such constant to its metadata device for the retrace only;
+    yields a map from the id of each copy to its CPU tensor, for
+    _restore_offloaded_constants.
+    """
+    swapped = {}
+    for node in gm.graph.nodes:
+        target = str(node.target)
+        if node.op != "get_attr" or not hasattr(gm, target):
+            continue
+        real, val = getattr(gm, target), node.meta.get("val", None)
+        if (
+            isinstance(real, torch.Tensor)
+            and isinstance(val, torch.Tensor)
+            and real.device != val.device
+        ):
+            swapped[target] = (real, real.to(val.device))
+    for target, (_, copy) in swapped.items():
+        delattr(gm, target)
+        setattr(gm, target, copy)
+    try:
+        yield {id(copy): real for real, copy in swapped.values()}
+    finally:
+        for target, (real, _) in swapped.items():
+            delattr(gm, target)
+            setattr(gm, target, real)
+
+
+def _restore_offloaded_constants(
+    gm: GraphModule, offloaded: dict[int, torch.Tensor]
+) -> None:
+    """Put the CPU tensors back wherever the retraced module kept a device copy."""
+    for node in gm.graph.nodes:
+        if node.op != "get_attr":
+            continue
+        real = offloaded.get(id(getattr(gm, str(node.target), None)))
+        if real is not None:
+            setattr(gm, str(node.target), real)
+
+
 def _fake_flat_args(gm: GraphModule) -> list[torch.Tensor]:
     """Build the example inputs the upstream retrace (make_fx) needs.
 
diff --git a/tests/py/dynamo/lowering/test_constant_folding.py b/tests/py/dynamo/lowering/test_constant_folding.py
index 5648932eea9..122e58f3c50 100644
--- a/tests/py/dynamo/lowering/test_constant_folding.py
+++ b/tests/py/dynamo/lowering/test_constant_folding.py
@@ -3,11 +3,13 @@
 
 import copy
 import unittest
+from unittest import mock
 
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.lowering.passes import complex_decomposition_adapter
 from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold
 
 
@@ -154,6 +156,52 @@ def test_aliased_view_models_compile_end_to_end(self):
 
                 torch.testing.assert_close(actual, expected)
 
+    @unittest.skipIf(
+        not torch.cuda.is_available()
+        or not complex_decomposition_adapter.has_complex_decomposition(),
+        "needs CUDA and the upstream complex decomposition",
+    )
+    def test_offloaded_constants_reach_the_complex_decomposition(self):
+        # offload_module_to_cpu leaves folded constants on the CPU; the complex
+        # decomposition retraces the graph and must not fall back over them.
+        class Rotary(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                angles = torch.outer(torch.arange(16.0), torch.ones(4))
+                self.register_buffer(
+                    "freqs", torch.polar(torch.ones_like(angles), angles)
+                )
+                # Read only by constant folding, so the retraced module drops it.
+                self.scale = torch.nn.Parameter(torch.ones(4))
+
+            def forward(self, x):
+                freqs = self.freqs[: x.shape[0]] * self.scale.mean()
+                z = torch.view_as_complex(x.reshape(x.shape[0], 4, 2))
+                return torch.view_as_real(z * freqs).flatten(1)
+
+        model = Rotary().eval().cuda()
+        example_input = torch.randn(8, 8, device="cuda")
+        with (
+            torch.no_grad(),
+            mock.patch.object(
+                complex_decomposition_adapter,
+                "complex_graph_detection",
+                side_effect=AssertionError("fell back to the legacy complex pass"),
+            ),
+        ):
+            expected = model(example_input)
+            exported_program = torch.export.export(model, (example_input,))
+            compiled_model = torch_tensorrt.dynamo.compile(
+                exported_program,
+                inputs=[example_input],
+                min_block_size=1,
+                offload_module_to_cpu=True,
+            )
+            actual = compiled_model(example_input)
+
+        torch.testing.assert_close(actual, expected)
+        self.assertEqual(model.scale.device.type, "cpu")
+
 
 if __name__ == "__main__":
     run_tests()

From 7b0a9e83587964796d9765d1a4d06c32a1a3c3b7 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 12:53:46 -0700
Subject: [PATCH 06/36] fix(conversion): convert float64 casts when
 truncate_double is set

With truncate_double=True, TensorRT runs float64 work in float32, and
the float64 ops around a cast already convert. The cast to float64
itself did not: the _to_copy validator only accepts dtypes TensorRT has,
so every x.to(torch.float64) stayed in PyTorch and cut the graph. A
rotary embedding that rotates in float64 in every attention block, as
Wan 2.2's rope_apply does, came out split into many engines, with a
PyTorch cast between each pair.

The validator now accepts a float64 cast of a 16, 32 or 64-bit float
when truncate_double is set, and impl.cast.to_copy casts it to float32
instead, the same way the cumsum and cat converters already treat a
float64 request. Other float64 requests that reach impl.cast.to_copy,
such as a copy into a float64 buffer, get the same treatment when the
source is a float, bool or 8-bit integer, which float32 holds exactly,
while a wider integer source still fails as before, because float32
cannot hold large integers exactly. _to_copy casts to float64 from any
other dtype still stay in PyTorch, as before; TensorRT cannot cast some
of them, such as float8, at all. So do casts whose PyTorch result is not
contiguous (channels_last, or preserving the layout of a non-contiguous
input), because the TensorRT copy is always contiguous. With dynamic
shapes, a cast that names a memory_format stays in PyTorch too, even in
the middle of a graph: the validator cannot know where partitioning will
put it, and if it ends up alone on an engine input, the placeholder copy
converter refuses an explicit memory_format and the build would fail.

This relies on #4898, which gives every engine output position its own
tensor. Without it, an engine that returns both a tensor and its float64
cast gets one binding for the two outputs, and the caller silently
receives wrong values in the wrong shape.
---
 .../dynamo/conversion/aten_ops_converters.py  |  46 ++++++-
 .../dynamo/conversion/impl/cast.py            |  10 ++
 tests/py/dynamo/conversion/test_casts.py      | 122 ++++++++++++++++++
 3 files changed, 176 insertions(+), 2 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 7cda4e4dc7a..4654518b90e 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2142,12 +2142,45 @@ def aten_ops_permute(
     )
 
 
+def _truncatable_float64_copy(node: Node) -> bool:
+    """Whether a float64 copy can run as a float32 copy and give PyTorch's result.
+
+    The source has to be a float type TensorRT casts to float32, and the PyTorch
+    result has to have the contiguous layout a TensorRT copy produces. The check
+    does not depend on where the copy sits, because partitioning can turn any copy
+    into an engine output and the converter is validated again then.
+    """
+    input_node = node.args[0] if node.args else None
+    input_meta = input_node.meta.get("val") if isinstance(input_node, Node) else None
+    if not isinstance(input_meta, torch.Tensor) or input_meta.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ):
+        return False
+    # Alone on a placeholder after partitioning, a dynamic copy must also pass
+    # _dynamic_placeholder_copy_supported, which refuses any explicit memory_format.
+    if node.kwargs.get("memory_format") is not None and node_has_dynamic_shapes(node):
+        return False
+    output_meta = node.meta.get("val")
+    if not isinstance(output_meta, torch.Tensor):
+        return False
+    # A data-dependent stride cannot be decided here, so leave the copy to PyTorch.
+    try:
+        return bool(output_meta.is_contiguous())
+    except Exception:
+        return False
+
+
 def to_copy_dtype_validator(
     placeholder_only: bool, settings: Optional[CompilationSettings] = None
 ) -> Callable[[Node, CompilationSettings], bool]:
     """Return validator for to_copy node with placeholder restrictions"""
 
-    def validate_dtype(to_copy_node: Node) -> bool:
+    def validate_dtype(
+        to_copy_node: Node, settings: Optional[CompilationSettings]
+    ) -> bool:
         """Returns true if the to_copy node can be converted to TRT
 
         Based on data type being casted to
@@ -2187,6 +2220,15 @@ def validate_dtype(to_copy_node: Node) -> bool:
         if "dtype" in to_copy_node.kwargs:
             if to_copy_node.kwargs["dtype"] in allowed_casts:
                 return True
+            # truncate_double means the caller accepts float32 for float64, and
+            # the converter maps it.
+            elif (
+                to_copy_node.kwargs["dtype"] is torch.float64
+                and settings is not None
+                and settings.truncate_double
+                and _truncatable_float64_copy(to_copy_node)
+            ):
+                return True
             else:
                 _LOGGER.debug(
                     f"_to_copy converter rejected node {to_copy_node} with dtype {to_copy_node.kwargs['dtype']}"
@@ -2206,7 +2248,7 @@ def validator(
         """
         # The placeholder restriction is satsfied if placeholder_only is the same
         # truth value as is_only_operator_on_placeholder(to_copy_node)
-        return validate_dtype(to_copy_node) and (
+        return validate_dtype(to_copy_node, settings) and (
             (not placeholder_only) ^ is_only_operator_on_placeholder(to_copy_node)
         )
 
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/cast.py b/py/torch_tensorrt/dynamo/conversion/impl/cast.py
index f4953f8ce91..1f373348007 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/cast.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/cast.py
@@ -34,6 +34,16 @@ def to_copy(
     if not isinstance(input, TRTTensor):
         input = get_trt_tensor(ctx, input, f"{name}_copy_tensor")
 
+    # TensorRT has no float64. With truncate_double the caller accepts float32,
+    # which holds these source types exactly but not every larger integer.
+    if (
+        ctx.compilation_settings.truncate_double
+        and _enums.dtype._from(dtype) == _enums.dtype.f64
+        and input.dtype
+        in {trt.bool, trt.int8, trt.uint8, trt.float16, trt.bfloat16, trt.float32}
+    ):
+        dtype = torch.float32
+
     # If cast is forced, insert identity layer regardless of whether the dtype
     # doesn't change
     if force_layer:
diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py
index 90d3e6d678a..ac2926ee2cb 100644
--- a/tests/py/dynamo/conversion/test_casts.py
+++ b/tests/py/dynamo/conversion/test_casts.py
@@ -5,7 +5,9 @@
 
 import torch
 import torch.nn as nn
+import torch_tensorrt
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo.conversion.aten_ops_converters import (
     to_copy_dtype_validator,
 )
@@ -94,6 +96,74 @@ def forward(self, x):
             inputs,
         )
 
+    def test_to_copy_float64_with_truncate_double(self):
+        class ToCopyDouble(nn.Module):
+            def forward(self, x):
+                y = x * 2
+                return y, torch.ops.aten._to_copy.default(y, dtype=torch.float64)
+
+        # aten.copy reaches the same cast without the _to_copy validator.
+        class CopyIntoDouble(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("dest", torch.zeros(1, 3, 10, dtype=torch.double))
+
+            def forward(self, x):
+                return torch.ops.aten.copy.default(self.dest, x * 2)
+
+        inputs = [torch.rand((1, 3, 10), dtype=torch.float)]
+        for model in (ToCopyDouble(), CopyIntoDouble()):
+            with self.subTest(type(model).__name__):
+                self.run_test(
+                    model,
+                    inputs,
+                    check_dtype=False,
+                    use_dynamo_tracer=True,
+                    require_full_compilation=True,
+                )
+
+    def test_to_copy_float64_dynamic_with_memory_format_after_pytorch_op(self):
+        # lgamma stays in PyTorch, so the cast ends up alone on an engine input.
+        class LgammaDouble(nn.Module):
+            def __init__(self, memory_format):
+                super().__init__()
+                self.memory_format = memory_format
+
+            def forward(self, x):
+                return torch.lgamma(x).to(
+                    torch.float64, memory_format=self.memory_format
+                )
+
+        for memory_format in (None, torch.contiguous_format, torch.preserve_format):
+            with self.subTest(memory_format=memory_format):
+                model = LgammaDouble(memory_format).eval().cuda()
+                x = torch.rand(3, 5, device="cuda") + 0.5
+                batch = torch.export.Dim("batch", min=1, max=8)
+                exported = torch.export.export(
+                    model, (x,), dynamic_shapes=({0: batch},)
+                )
+                compiled = torch_tensorrt.dynamo.compile(
+                    exported,
+                    arg_inputs=[
+                        torch_tensorrt.Input(
+                            min_shape=(1, 5),
+                            opt_shape=(3, 5),
+                            max_shape=(8, 5),
+                            dtype=torch.float32,
+                        )
+                    ],
+                    min_block_size=1,
+                    truncate_double=True,
+                )
+                # Only the cast without a memory_format goes to TensorRT.
+                self.assertEqual(
+                    sum("_run_on_acc" in name for name, _ in compiled.named_children()),
+                    int(memory_format is None),
+                )
+                for n in (1, 3, 8):
+                    x = torch.rand(n, 5, device="cuda") + 0.5
+                    torch.testing.assert_close(compiled(x), model(x))
+
     def test_to_copy_multiple_returns(self):
         class ToCopyReturns(nn.Module):
             def forward(self, x):
@@ -133,6 +203,58 @@ def to_copy_node(device):
         self.assertFalse(validator(to_copy_node("cpu")))
         self.assertTrue(validator(to_copy_node("cuda")))
 
+    def test_to_copy_validator_float64_keeps_layout_in_pytorch(self):
+        class DoubleCopy(nn.Module):
+            def __init__(self, memory_format, view):
+                super().__init__()
+                self.memory_format = memory_format
+                self.view = view
+
+            def forward(self, x):
+                y = self.view(x + 1)
+                return torch.ops.aten._to_copy.default(
+                    y, dtype=torch.float64, memory_format=self.memory_format
+                )
+
+        def to_copy_node(memory_format, dtype=torch.float32, view=lambda y: y):
+            x = torch.ones(2, 3, 4, 5, dtype=dtype)
+            exported = torch.export.export(DoubleCopy(memory_format, view), (x,))
+            return next(
+                node
+                for node in exported.graph.nodes
+                if node.target is torch.ops.aten._to_copy.default
+            )
+
+        validator = to_copy_dtype_validator(placeholder_only=False)
+        settings = CompilationSettings(truncate_double=True)
+        self.assertTrue(validator(to_copy_node(torch.contiguous_format), settings))
+        self.assertFalse(validator(to_copy_node(torch.channels_last), settings))
+        # Copying a strided input gives a contiguous result, which TensorRT matches.
+        self.assertTrue(
+            validator(
+                to_copy_node(torch.preserve_format, view=lambda y: y[..., ::2]),
+                settings,
+            )
+        )
+        self.assertFalse(
+            validator(
+                to_copy_node(torch.preserve_format, view=lambda y: y.transpose(1, 2)),
+                settings,
+            )
+        )
+        # TensorRT cannot cast float8 to float32.
+        self.assertFalse(
+            validator(
+                to_copy_node(torch.contiguous_format, torch.float8_e4m3fn), settings
+            )
+        )
+        self.assertFalse(
+            validator(
+                to_copy_node(torch.contiguous_format),
+                CompilationSettings(truncate_double=False),
+            )
+        )
+
 
 if __name__ == "__main__":
     run_tests()

From ac2d7c0ff3a64be4ca1fbdc9fc1d6fd46b6a0784 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Wed, 28 Oct 2026 02:42:24 -0700
Subject: [PATCH 07/36] fix(lowering): keep attention-mask fold exclusions
 through retraces

When attention is decomposed, the decomposition marks the aranges behind
the attention mask so that constant folding leaves them alone. The mark is
node metadata. When a graph has complex operations, the complex
decomposition adapter retraces the whole graph and then folds constants
again. The retrace builds new nodes without that metadata, and the
graph-level rule cannot find the pattern again because the attention op
is already gone. So the aranges are folded after all, and a decomposed
1024x1024 causal mask becomes an engine constant (4 MB in float32)
instead of being computed. A second run_decompositions, for example on a
program the user already decomposed, loses the marks the same way.

Tag the aranges in node.meta["custom"] instead, which retraces copy to
the new nodes, and let the attention_mask_arange rule select tagged
nodes. Disabling the rule still makes them foldable.
---
 .../dynamo/lowering/_decompositions.py        |  2 +-
 .../constant_fold_exclusions/_core.py         |  3 +-
 .../attention_mask.py                         | 19 +++---
 .../passes/mark_constant_fold_exclusions.py   | 10 +--
 ..._attention_mask_constant_fold_exclusion.py | 61 ++++++++++++++++---
 5 files changed, 71 insertions(+), 24 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
index 56cf0ca052f..14d12d2ccb1 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
@@ -509,7 +509,7 @@ def scaled_dot_product_attention_decomposition(
         attn_bias = torch.where(temp_mask.logical_not(), negative, zero)
 
     if attn_mask is not None:
-        # Unconditional: mark_constant_fold_exclusions revokes these marks when the
+        # Unconditional: mark_constant_fold_exclusions ignores the tags when the
         # rule is disabled, so the decompositions do not need the setting.
         exclude_attn_mask_aranges_from_constant_fold(attn_mask)
         if attn_mask.dtype == torch.bool:
diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py
index c9a3c4fba2a..9b3a96a69d7 100644
--- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py
@@ -15,8 +15,7 @@ def _mark_constant_fold_exclusion(nodes: Iterable[torch.fx.Node], rule_id: str)
     """Record which rule wants ``nodes`` kept out of constant folding.
 
     The marks carry their rule ID rather than a bare flag so the exclusion pass
-    can revoke the ones belonging to disabled rules, whichever marking path
-    produced them.
+    can revoke the ones belonging to disabled rules.
     """
     for node in nodes:
         node.meta.setdefault(CONSTANT_FOLD_EXCLUSION_META_KEY, set()).add(rule_id)
diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py
index 7abe4c753ca..d60234d431b 100644
--- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py
@@ -7,12 +7,10 @@
 from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor
 from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot
 
-from ._core import (
-    _mark_constant_fold_exclusion,
-    register_constant_fold_exclusion_rule,
-)
+from ._core import register_constant_fold_exclusion_rule
 
 ATTENTION_MASK_ARANGE_RULE_ID = "attention_mask_arange"
+_ATTENTION_MASK_ARANGE_TAG = "_torch_tensorrt_attention_mask_arange"
 
 
 def _is_arange_node(node: torch.fx.Node) -> bool:
@@ -62,7 +60,7 @@ def _find_attention_mask_aranges(
 
 
 def exclude_attn_mask_aranges_from_constant_fold(attn_mask: torch.Tensor) -> None:
-    """Mark aranges behind an attention mask while tracing a decomposition.
+    """Tag aranges behind an attention mask while tracing a decomposition.
 
     ``run_decompositions`` invokes decompositions once for functionalization and
     again while proxy tracing. Only the proxy-tracing invocation has an FX graph
@@ -82,15 +80,18 @@ def exclude_attn_mask_aranges_from_constant_fold(attn_mask: torch.Tensor) -> Non
     mask_node = getattr(proxy, "node", None)
 
     if isinstance(mask_node, torch.fx.Node):
-        _mark_constant_fold_exclusion(
-            _find_attention_mask_aranges(mask_node),
-            ATTENTION_MASK_ARANGE_RULE_ID,
-        )
+        # Retraces keep node.meta["custom"], so the rule below still finds
+        # these aranges after the attention op has been decomposed away.
+        for node in _find_attention_mask_aranges(mask_node):
+            node.meta.setdefault("custom", {})[_ATTENTION_MASK_ARANGE_TAG] = True
 
 
 @register_constant_fold_exclusion_rule(ATTENTION_MASK_ARANGE_RULE_ID)
 def _attention_mask_arange_rule(node: torch.fx.Node) -> Iterable[torch.fx.Node]:
     """Select aranges feeding an attention mask."""
+    if node.meta.get("custom", {}).get(_ATTENTION_MASK_ARANGE_TAG):
+        return (node,)
+
     # Every SDPA overload that takes a mask keeps it at positional index 3.
     # _scaled_dot_product_flash_attention is absent because it has no mask.
     attention_mask_args: dict[Any, tuple[int, str]] = {
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
index 5c1150afdf3..311d23fe5e9 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
@@ -18,11 +18,11 @@ def mark_constant_fold_exclusions(
     """Apply the registered rules that exclude FX nodes from constant folding.
 
     This pass is the single authority on which rules are in effect. It runs
-    immediately before ``constant_fold`` and is the only marking path that sees
-    ``settings``: rules that mark nodes while a decomposition is traced run
-    during ``run_decompositions``, long before a settings object is reachable.
-    Those marks are therefore revoked here rather than suppressed where they are
-    made, so a caller only has to communicate the disabled rules once.
+    immediately before ``constant_fold`` and is the only code that marks nodes:
+    decompositions run during ``run_decompositions``, long before a settings
+    object is reachable, so they only tag nodes for a rule to select here. As a
+    guard, it also revokes marks of disabled rules that a node already carries,
+    so a caller only has to communicate the disabled rules once.
     """
     disabled_rule_ids = validate_disabled_constant_fold_exclusions(
         settings.disabled_constant_fold_exclusions if settings is not None else ()
diff --git a/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py
index 66b0d6a1f0e..6b19c929ba9 100644
--- a/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py
+++ b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py
@@ -46,11 +46,16 @@ def _export(self):
         return torch.export.export(self.AttentionWithCausalMask(), inputs)
 
     def _assert_only_attention_aranges_survive(
-        self, decompose_attention, disabled_constant_fold_exclusions=()
+        self,
+        decompose_attention,
+        disabled_constant_fold_exclusions=(),
+        decompositions_runs=1,
     ):
-        exported_program = self._export().run_decompositions(
-            get_decompositions(decompose_attention=decompose_attention)
-        )
+        exported_program = self._export()
+        for _ in range(decompositions_runs):
+            exported_program = exported_program.run_decompositions(
+                get_decompositions(decompose_attention=decompose_attention)
+            )
         gm = post_lowering(
             exported_program.module(),
             CompilationSettings(
@@ -78,6 +83,48 @@ def test_decomposed_attention_mask_aranges_are_not_folded(self):
     def test_native_attention_mask_aranges_are_not_folded(self):
         self._assert_only_attention_aranges_survive(decompose_attention=False)
 
+    def test_decomposed_attention_mask_aranges_survive_a_second_decomposition(self):
+        self._assert_only_attention_aranges_survive(
+            decompose_attention=True, decompositions_runs=2
+        )
+
+    def test_decomposed_attention_mask_aranges_survive_complex_retrace(self):
+        class AttentionWithComplexOp(self.AttentionWithCausalMask):
+            def forward(self, query, key, value, attention_mask, rotation):
+                attention, unrelated_arange = super().forward(
+                    query, key, value, attention_mask
+                )
+                rotation = torch.view_as_complex(rotation)
+                return attention, unrelated_arange, torch.view_as_real(rotation**2)
+
+        inputs = (
+            torch.randn(1, 2, 8, 16, device="cuda"),
+            torch.randn(1, 2, 8, 16, device="cuda"),
+            torch.randn(1, 2, 8, 16, device="cuda"),
+            torch.ones(8, 8, dtype=torch.bool, device="cuda"),
+            torch.randn(4, 2, device="cuda"),
+        )
+        exported_program = torch.export.export(
+            AttentionWithComplexOp(), inputs
+        ).run_decompositions(get_decompositions(decompose_attention=True))
+        gm = post_lowering(
+            exported_program.module(),
+            CompilationSettings(use_complex_decomposition=True),
+        )
+        arange_nodes = [
+            node
+            for node in gm.graph.nodes
+            if node.op == "call_function"
+            and getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange
+        ]
+        self.assertEqual(len(arange_nodes), 2)
+        self.assertTrue(
+            all(
+                node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY, False)
+                for node in arange_nodes
+            )
+        )
+
     def test_native_attention_rules_can_be_disabled(self):
         exported_program = self._export().run_decompositions(
             get_decompositions(decompose_attention=False)
@@ -100,9 +147,9 @@ def test_native_attention_rules_can_be_disabled(self):
     def test_decomposed_attention_rules_can_be_disabled(self):
         """post_lowering is the only place a rule has to be disabled.
 
-        The decompositions mark unconditionally while tracing, well before any
-        settings object is reachable, so post_lowering revokes those marks
-        instead of the caller having to communicate the setting twice.
+        The decompositions tag unconditionally while tracing, well before any
+        settings object is reachable, and post_lowering skips the disabled rule
+        that would turn the tags into marks, so the caller sets it only once.
         """
         exported_program = self._export().run_decompositions(
             get_decompositions(decompose_attention=True)

From d3741c2732b3a114c9bf6b2f90d78bd964b051ef Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Wed, 28 Oct 2026 07:55:51 -0700
Subject: [PATCH 08/36] fix(runtime): keep shape-tensor inputs alive until the
 engine runs

The Python runtime copies each shape-tensor input into a NumPy array and
hands TensorRT the array's address. The array was a loop variable, so it
was freed as soon as the next input was handled, before the engine ran.
A later allocation could reuse that memory: with three shape inputs, the
third array often lands where the first was, and TensorRT then reads the
third value for the first input. The engine fails with a shape error or
runs with the wrong shapes. torch.compile with dynamic shapes passes
symbolic sizes as shape inputs, so it is the usual way to hit this.

Keep the arrays in a list on the engine until the next call sets up its
inputs.
---
 .../dynamo/runtime/_TRTEngine.py              |  6 ++++
 .../dynamo/runtime/test_host_shape_inputs.py  | 35 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py
index e66b75a3a7c..64354b0cbef 100644
--- a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py
+++ b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py
@@ -266,6 +266,7 @@ def __init__(
         self._input_buffers: List[torch.Tensor] = []
         self._output_buffers: List[torch.Tensor] = []
         self._bound_inputs_by_name: Dict[str, torch.Tensor] = {}
+        self._shape_input_buffers: List[Any] = []
         self._caller_stream: Optional[torch.cuda.Stream] = None
         self._engine_stream: Optional[torch.cuda.Stream] = None
         self._owned_pool_stream: Optional[torch.cuda.Stream] = None
@@ -407,6 +408,7 @@ def __setstate__(self, state: Any) -> None:
         self._input_buffers = []
         self._output_buffers = []
         self._bound_inputs_by_name = {}
+        self._shape_input_buffers = []
         self._caller_stream = None
         self._engine_stream = None
         self._owned_pool_stream = None
@@ -1192,6 +1194,7 @@ def setup_input_tensors(
         # create_output_tensors / the output-binding loop to alias outputs to
         # their source-input storage.
         self._bound_inputs_by_name = {}
+        self._shape_input_buffers = []
         for i, binding in enumerate(self._input_binding_infos):
             input_name = binding.name
 
@@ -1208,6 +1211,9 @@ def setup_input_tensors(
 
             if binding.is_shape_tensor:
                 inputs_cpu = contiguous_inputs[i].cpu().to(torch.int64).numpy().copy()
+                # TensorRT reads this host memory when the engine runs, so the
+                # array must outlive this loop.
+                self._shape_input_buffers.append(inputs_cpu)
                 self.context.set_tensor_address(input_name, inputs_cpu.ctypes.data)
             else:
                 self.context.set_input_shape(
diff --git a/tests/py/dynamo/runtime/test_host_shape_inputs.py b/tests/py/dynamo/runtime/test_host_shape_inputs.py
index f607ebba86f..49d312a18fb 100644
--- a/tests/py/dynamo/runtime/test_host_shape_inputs.py
+++ b/tests/py/dynamo/runtime/test_host_shape_inputs.py
@@ -3,6 +3,7 @@
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo.runtime import TorchTensorRTModule
+from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine
 
 
 class ShapeInputModel(nn.Module):
@@ -69,6 +70,40 @@ def test_host_and_device_shape_input_agree(self):
         device_out = self._run_engine(trt_mod, x, 64, "cuda")
         torch.testing.assert_close(host_out, device_out, rtol=1e-5, atol=1e-5)
 
+    def test_python_runtime_keeps_every_shape_input(self):
+        class ThreeShapeInputModel(nn.Module):
+            def forward(self, x, a, b, c):
+                ranges = [torch.arange(0, n, 1, device=x.device) for n in (a, b, c)]
+                return torch.cat(ranges).to(torch.float32) + x.sum()
+
+        model = ThreeShapeInputModel().eval().cuda()
+        inputs = (
+            torch.randn(8, device="cuda"),
+            *(torch.tensor(n, dtype=torch.int64, device="cuda") for n in (3, 5, 7)),
+        )
+        gm = torch_tensorrt.dynamo.compile(
+            torch.export.export(model, inputs),
+            inputs=list(inputs),
+            min_block_size=1,
+            pass_through_build_failures=True,
+            assume_dynamic_shape_support=True,
+        )
+        trt_mod = next(
+            m for _, m in gm.named_children() if isinstance(m, TorchTensorRTModule)
+        )
+        # A build with the C++ runtime backs the module with the C++ engine, so
+        # drive the Python engine directly from the same serialized info.
+        engine = TRTEngine(trt_mod._pack_engine_info())
+        shape_names = sorted(
+            (name for name in engine.in_binding_names if name != "x"),
+            key=lambda name: (len(name), name),
+        )
+        args = {"x": inputs[0], **dict(zip(shape_names, inputs[1:]))}
+        output = engine.execute([args[name] for name in engine.in_binding_names])
+        if isinstance(output, (list, tuple)):
+            (output,) = output
+        torch.testing.assert_close(output, model(*inputs))
+
 
 if __name__ == "__main__":
     run_tests()

From bd82800152b1ffe917227898f3c8716b72f26fc4 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 27 Oct 2026 23:36:12 -0700
Subject: [PATCH 09/36] fix(lowering): decompose integer scatter_add without a
 loop over the scattered extent

The scatter_add decomposition loops over every slice of src along dim and
emits a select, unsqueeze, zeros constant, scatter and add for each one,
so a scatter_add over n slices becomes about 8n graph nodes and n
constants the size of the input. A mixture-of-experts router that counts
tokens per expert with an int64 scatter_add grows by about 290,000 nodes,
and the TensorRT build then takes hours.

TensorRT scatter overwrites colliding indices, which is why the loop
scatters one slice at a time. For int32 and int64, compare every index
with every output position along dim and sum the matching src values
instead. Integer addition gives the same result in any order, so the
output is exact. As with TensorRT gather and scatter, indices must be in
range: an out-of-range index is ignored rather than raising. The graph no
longer grows with n. Only the part of src covered by index is used, as in
PyTorch. The sum multiplies by the 0/1 match mask because the where
converter narrows int64 to int32. The compare adds one axis, so an index
of rank 8, the TensorRT limit, keeps the loop. The comparison holds
index.numel() * input.shape[dim] elements, so above 2^26 of them, well
under TensorRT's tensor volume limit, the loop is kept too, as it is when
the bound cannot be proven, for example for a dynamic dim with no upper
bound; those cases keep the large graphs described above. The loop now
runs over the extent of index along dim, which is all PyTorch reads, so a
longer src no longer raises an IndexError, and filter_decomposition_table
checks that extent.

Every other dtype keeps the loop decomposition. Floating point sums
would round differently in another order, bool scatter_add is a logical
or rather than a count, and the converters cannot compute the narrow
integer types. When index has a dynamic size along dim, the op stays in
PyTorch or the loop specializes it.

With a constant index, constant folding would turn the whole comparison
into one engine constant. A new constant-fold exclusion rule,
scatter_add_positions, keeps the positions arange out of folding, like
attention_mask_arange does for attention masks. The decomposition tags
the arange in node.meta["custom"], which FX retraces keep, so the rule
still finds it after the complex decomposition retrace or a second
run_decompositions. The rule applies only when src varies at runtime (it
reads an input or a random op): with constant src and index, the update
sum still folds to an output-sized constant. For a 128x64 constant index
into a 4096x64 int64 output, the engine is about 90 KB, compared with
about 3.2 MB for the loop. A few updates into a very wide output can run
slower: 16 rows into a 262144x8 output took 0.47 ms instead of 0.36 ms.
---
 py/torch_tensorrt/dynamo/_compiler.py         |  15 ++-
 py/torch_tensorrt/dynamo/_settings.py         |   5 +-
 .../dynamo/lowering/_decomposition_filter.py  |  13 ++-
 .../dynamo/lowering/_decompositions.py        |  79 +++++++++++++-
 .../constant_fold_exclusions/__init__.py      |   2 +
 .../constant_fold_exclusions/scatter_add.py   |  77 +++++++++++++
 .../py/dynamo/lowering/test_decompositions.py | 103 ++++++++++++++++--
 7 files changed, 266 insertions(+), 28 deletions(-)
 create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index 3bbcf33eed3..8e3bdaa5cd3 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -225,8 +225,9 @@ def cross_compile_for_windows(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
@@ -641,8 +642,9 @@ def compile(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
@@ -2110,8 +2112,9 @@ def convert_exported_program_to_serialized_trt_engine(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py
index 206f5354cee..1f7fabfbd0d 100644
--- a/py/torch_tensorrt/dynamo/_settings.py
+++ b/py/torch_tensorrt/dynamo/_settings.py
@@ -159,8 +159,9 @@ class CompilationSettings:
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
diff --git a/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py b/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py
index 191c905b6a9..415311d3f2c 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py
@@ -16,13 +16,16 @@ def _has_symbolic_scatter_add_extent(graph_module: torch.fx.GraphModule) -> bool
         if node.target != torch.ops.aten.scatter_add.default or len(node.args) < 4:
             continue
         dim = node.args[1]
-        src_node = node.args[3]
-        if not isinstance(dim, int) or not isinstance(src_node, torch.fx.Node):
+        # The decomposition loops over the extent of index along dim.
+        index_node = node.args[2]
+        if not isinstance(dim, int) or not isinstance(index_node, torch.fx.Node):
             continue
-        src_val = src_node.meta.get("val", src_node.meta.get("example_value"))
-        if not isinstance(src_val, torch.Tensor) or not src_val.ndim:
+        index_val = index_node.meta.get("val", index_node.meta.get("example_value"))
+        if not isinstance(index_val, torch.Tensor) or not index_val.ndim:
             continue
-        if isinstance(src_val.shape[get_positive_dim(dim, src_val.ndim)], torch.SymInt):
+        if isinstance(
+            index_val.shape[get_positive_dim(dim, index_val.ndim)], torch.SymInt
+        ):
             return True
 
     return False
diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
index 14d12d2ccb1..c75fa9e998e 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
@@ -13,6 +13,7 @@
     _get_decomp_for_cia,
 )
 from torch._ops import OpOverload
+from torch.fx.experimental.symbolic_shapes import statically_known_true
 from torch_tensorrt.dynamo._defaults import default_device
 from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim
 from torch_tensorrt.dynamo.utils import to_torch_device
@@ -28,6 +29,9 @@
 from .constant_fold_exclusions.attention_mask import (
     exclude_attn_mask_aranges_from_constant_fold,
 )
+from .constant_fold_exclusions.scatter_add import (
+    exclude_scatter_add_positions_from_constant_fold,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -292,10 +296,12 @@ def scatter_add_decomposition(
     index: torch.Tensor,
     src_tensor: torch.Tensor,
 ) -> torch.Tensor:
+    if _use_integer_scatter_add(input_tensor, dim, index, src_tensor):
+        return _scatter_add_integer_decomposition(input_tensor, dim, index, src_tensor)
+
     scatter_add_tensor = input_tensor
-    src_shape = list(src_tensor.shape)
-    src_dim = src_shape[dim]
-    for i in range(0, src_dim):
+    # PyTorch reads src only where index has values, so loop over index.
+    for i in range(0, index.shape[dim]):
         to_scatter_tensor = torch.zeros(input_tensor.shape, dtype=input_tensor.dtype)
         # index and src slice
         src_slice = torch.select(src_tensor, dim, i)
@@ -320,6 +326,73 @@ def scatter_add_decomposition(
     return scatter_add_tensor
 
 
+# The comparison holds index.numel() * input.shape[dim] elements. Keep it well
+# under TensorRT's tensor volume limit and small enough (512 MiB as int64) for
+# constant folding or a PyTorch fallback to materialize.
+_MAX_INTEGER_SCATTER_ADD_COMPARISONS = 1 << 26
+
+
+def _use_integer_scatter_add(
+    input_tensor: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+    src_tensor: torch.Tensor,
+) -> bool:
+    # Integer reductions can regroup updates without changing their values.
+    if src_tensor.dtype not in (torch.int32, torch.int64):
+        return False
+    # The comparison adds one axis, and TensorRT tensors have at most 8.
+    if index.dim() >= 8:
+        return False
+    if input_tensor.dim() == 0:
+        extent = 1
+    elif index.dim() and isinstance(index.shape[dim], int):
+        extent = input_tensor.shape[dim]
+    else:
+        # Otherwise keep the loop, as before. It specializes a dynamic extent.
+        return False
+    comparisons = index.numel() * extent
+    return statically_known_true(comparisons <= _MAX_INTEGER_SCATTER_ADD_COMPARISONS)
+
+
+def _scatter_add_integer_decomposition(
+    input_tensor: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+    src_tensor: torch.Tensor,
+) -> torch.Tensor:
+    # TensorRT scatter overwrites colliding indices, so compare each index with
+    # every output position along dim and sum the matching src values.
+    if statically_known_true(index.numel() == 0):
+        return input_tensor.clone()
+    if input_tensor.dim() == 0:
+        # PyTorch treats a scalar destination as one element along dim 0.
+        return _scatter_add_integer_decomposition(
+            input_tensor.reshape(1), 0, index.reshape(-1), src_tensor.reshape(-1)
+        ).reshape(())
+    dim = get_positive_dim(dim, input_tensor.dim())
+    for k, size in enumerate(index.shape):
+        if not statically_known_true(src_tensor.shape[k] == size):
+            # A slice view would add guards on dynamic sizes, a gather does not.
+            src_tensor = src_tensor.index_select(
+                k, torch.arange(size, device=src_tensor.device)
+            )
+    index = index.movedim(dim, -1).unsqueeze(-2)
+    src_tensor = src_tensor.movedim(dim, -1).unsqueeze(-2)
+    positions = torch.arange(input_tensor.shape[dim], device=index.device)
+    exclude_scatter_add_positions_from_constant_fold(positions)
+    hits = index == positions.view(-1, 1)
+    # The where converter narrows int64 to int32, multiplying keeps the dtype.
+    added = (hits.to(src_tensor.dtype) * src_tensor).sum(-1, dtype=src_tensor.dtype)
+    added = added.movedim(-1, dim)
+    pad = []
+    for k in reversed(range(input_tensor.dim())):
+        pad += [0, input_tensor.shape[k] - added.shape[k]]
+    if not all(statically_known_true(p == 0) for p in pad):
+        added = torch.constant_pad_nd(added, pad, 0)
+    return input_tensor + added
+
+
 # enum class for reduce operation of scatter_reduce
 class ReduceOperation(Enum):
     SUM = ("Sum reduce operation", lambda x, y: torch.add(x, y))
diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py
index 8e88946240c..855faaf12d9 100644
--- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py
@@ -9,11 +9,13 @@
     validate_disabled_constant_fold_exclusions,
 )
 from .attention_mask import ATTENTION_MASK_ARANGE_RULE_ID
+from .scatter_add import SCATTER_ADD_POSITIONS_RULE_ID
 
 __all__ = [
     "ATTENTION_MASK_ARANGE_RULE_ID",
     "CONSTANT_FOLD_EXCLUSION_META_KEY",
     "ConstantFoldExclusionRule",
+    "SCATTER_ADD_POSITIONS_RULE_ID",
     "register_constant_fold_exclusion_rule",
     "unregister_constant_fold_exclusion_rule",
     "validate_disabled_constant_fold_exclusions",
diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
new file mode 100644
index 00000000000..a9fc83af230
--- /dev/null
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import Iterable
+
+import torch
+from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor
+from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot
+
+from ._core import register_constant_fold_exclusion_rule
+
+SCATTER_ADD_POSITIONS_RULE_ID = "scatter_add_positions"
+_POSITIONS_TAG = "_torch_tensorrt_scatter_add_positions"
+
+
+def exclude_scatter_add_positions_from_constant_fold(positions: torch.Tensor) -> None:
+    """Tag the positions arange of the scatter_add decomposition.
+
+    With a constant index, folding it would store the whole comparison of every
+    index with every position as an engine constant. Retraces keep
+    ``node.meta["custom"]``, so the rule below finds the tag after the
+    decomposed op is gone, for example after the complex decomposition retrace
+    or a second ``run_decompositions``.
+    """
+    proxy_mode = get_proxy_mode()
+    if proxy_mode is None:
+        return
+
+    tracked = get_proxy_slot(
+        mb_unwrap_functional_tensor(positions), proxy_mode.tracer, default=None
+    )
+    node = getattr(getattr(tracked, "proxy", tracked), "node", None)
+    if isinstance(node, torch.fx.Node):
+        node.meta.setdefault("custom", {})[_POSITIONS_TAG] = True
+
+
+def _varies_at_runtime(node: torch.fx.Node) -> bool:
+    """Whether constant folding cannot compute this node ahead of time."""
+    nodes_to_visit, visited = [node], set()
+    while nodes_to_visit:
+        current = nodes_to_visit.pop()
+        target_tags = getattr(current.target, "tags", ())
+        if (
+            current.op == "placeholder"
+            or torch.Tag.nondeterministic_seeded in target_tags
+        ):
+            return True
+        if current not in visited:
+            visited.add(current)
+            nodes_to_visit.extend(current.all_input_nodes)
+    return False
+
+
+@register_constant_fold_exclusion_rule(SCATTER_ADD_POSITIONS_RULE_ID)
+def _scatter_add_positions_rule(node: torch.fx.Node) -> Iterable[torch.fx.Node]:
+    """Select the tagged positions arange when src varies at runtime.
+
+    The decomposition multiplies ``(index == positions.view(-1, 1)).to(dtype)``
+    by src. With constant src, a constant index makes the whole sum constant,
+    and it folds to a result the size of the output, so the positions stay
+    foldable. A runtime index keeps the comparison live either way.
+    """
+    if not node.meta.get("custom", {}).get(_POSITIONS_TAG):
+        return ()
+    # Several scatter_adds can share the positions, for example after CSE.
+    frontier, comparison, sources = [node], {node}, []
+    for _ in range(4):  # view, eq, cast, mul
+        frontier = [user for current in frontier for user in current.users]
+        for user in frontier:
+            if user.target == torch.ops.aten.mul.Tensor:
+                sources += [
+                    arg for arg in user.all_input_nodes if arg not in comparison
+                ]
+        comparison.update(frontier)
+    if sources and not any(map(_varies_at_runtime, sources)):
+        return ()
+    return (node,)
diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py
index 629f88f0501..573a6f2939a 100644
--- a/tests/py/dynamo/lowering/test_decompositions.py
+++ b/tests/py/dynamo/lowering/test_decompositions.py
@@ -1362,17 +1362,14 @@ def forward(self, base, index, src):
         model = ScatterAdd().eval()
         base = torch.zeros((16, 8))
         index = torch.zeros((4, 8), dtype=torch.int64)
-        src = torch.ones((4, 8))
-        src_rows = torch.export.Dim("src_rows", min=1, max=16)
+        src = torch.ones((16, 8))
+        index_rows = torch.export.Dim("index_rows", min=1, max=16)
 
+        # The decomposition loops over index, so only its extent matters.
         dynamic_export = torch.export.export(
             model,
             (base, index, src),
-            dynamic_shapes={
-                "base": {},
-                "index": {0: src_rows},
-                "src": {0: src_rows},
-            },
+            dynamic_shapes={"base": {}, "index": {0: index_rows}, "src": {}},
         )
         base_decompositions = get_decompositions()
         dynamic_decompositions = filter_decomposition_table(
@@ -1395,9 +1392,9 @@ def forward(self, base, index, src):
             for node in dynamic_export.graph_module.graph.nodes
             if node.target == torch.ops.aten.scatter_add.default
         )
-        dynamic_src = dynamic_scatter.args[3]
-        self.assertIsInstance(dynamic_src, torch.fx.Node)
-        dynamic_src.meta["example_value"] = dynamic_src.meta.pop("val")
+        dynamic_index = dynamic_scatter.args[2]
+        self.assertIsInstance(dynamic_index, torch.fx.Node)
+        dynamic_index.meta["example_value"] = dynamic_index.meta.pop("val")
         dynamo_decompositions = filter_decomposition_table(
             get_decompositions(), dynamic_export.graph_module
         )
@@ -1409,6 +1406,75 @@ def forward(self, base, index, src):
         )
         self.assertIn(torch.ops.aten.scatter_add.default, static_decompositions)
 
+    @parameterized.expand(
+        [
+            ("rule_enabled", (), 1, "input", True),
+            ("rule_disabled", {"scatter_add_positions"}, 1, "input", False),
+            ("decomposed_twice", (), 2, "input", True),
+            ("constant_src_folds", (), 1, "ones", False),
+            ("random_src", (), 1, "randint", True),
+        ]
+    )
+    def test_scatter_add_constant_index_positions_not_folded(
+        self,
+        _,
+        disabled_constant_fold_exclusions,
+        decompositions_runs,
+        src_kind,
+        positions_kept,
+    ):
+        class ConstantIndexScatterAdd(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("index", torch.randint(0, 64, (16, 8)))
+
+            def forward(self, input, src):
+                if src_kind == "ones":
+                    src = torch.ones_like(self.index)
+                elif src_kind == "randint":
+                    src = torch.randint_like(self.index, 5)
+                return input.scatter_add(0, self.index, src)
+
+        exported = torch.export.export(
+            ConstantIndexScatterAdd().cuda(),
+            (
+                torch.zeros(64, 8, dtype=torch.int64).cuda(),
+                torch.ones(16, 8, dtype=torch.int64).cuda(),
+            ),
+        )
+        for _ in range(decompositions_runs):
+            exported = exported.run_decompositions(get_decompositions())
+        gm = post_lowering(
+            exported.module(),
+            CompilationSettings(
+                disabled_constant_fold_exclusions=disabled_constant_fold_exclusions
+            ),
+        )
+        has_arange = any(
+            getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange
+            for node in gm.graph.nodes
+        )
+        self.assertEqual(has_arange, positions_kept)
+
+    def test_scatter_add_large_comparison_keeps_loop(self):
+        class ScatterAdd(torch.nn.Module):
+            def forward(self, input, index, src):
+                return input.scatter_add(0, index, src)
+
+        def lowered_targets(updates):
+            inputs = (
+                torch.zeros(2**20, dtype=torch.int64),
+                torch.zeros(updates, dtype=torch.int64),
+                # src longer than index also exercises the loop above the bound.
+                torch.ones(updates + 1, dtype=torch.int64),
+            )
+            exported = torch.export.export(ScatterAdd(), inputs)
+            lowered = exported.run_decompositions(get_decompositions())
+            return {node.target for node in lowered.graph.nodes}
+
+        self.assertNotIn(torch.ops.aten.scatter.src, lowered_targets(64))
+        self.assertIn(torch.ops.aten.scatter.src, lowered_targets(128))
+
     @parameterized.expand(
         [
             (
@@ -1448,12 +1514,22 @@ def forward(self, base, index, src):
                 1,
                 torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1], [3, 2, 1, 2]]).cuda(),
                 torch.tensor(
-                    [[1, 2, 3, 1], [5, 6, 5, 5], [2, 4, 3, 2]], dtype=torch.int32
+                    [[1, 2, 3, 1], [5, 6, 5, 5], [2, 4, 3, 2]], dtype=torch.float32
                 ).cuda(),
                 {
                     torch.ops.aten.add.Tensor,
                 },
             ),
+            (
+                "scatter_add_zero_dim_src_larger_than_index",
+                0,
+                torch.tensor([[0, 1, 2, 0], [2, 2, 0, 1]]).cuda(),
+                torch.tensor(
+                    [[1, 2, 3, 4, 5], [6, 2**40, 8, 9, 10], [11, 12, 13, 14, 15]],
+                    dtype=torch.int64,
+                ).cuda(),
+                {torch.ops.aten.add.Tensor},
+            ),
         ]
     )
     def test_scatter_add(self, _, dim, index, src, expected_ops_param):
@@ -1467,8 +1543,11 @@ def forward(self, input):
         # Operations expected to be included in the traced graph after decompositions
         expected_ops = expected_ops_param
         unexpected_ops = {torch.ops.aten.scatter_add.default}
+        if not src.is_floating_point():
+            unexpected_ops.add(torch.ops.aten.scatter.src)
 
-        input = torch.zeros(3, 5, dtype=torch.int32).cuda()
+        # Nonzero, so cells without updates must keep their values.
+        input = torch.arange(15).reshape(3, 5).to(src.dtype).cuda()
         inputs = [input]
 
         fx_graph = torch.fx.symbolic_trace(TestModule())

From 43ac01f7d70a1cd3968bec6e542ad44a2987ff47 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Wed, 28 Oct 2026 22:22:21 -0700
Subject: [PATCH 10/36] fix(conversion): index_put with a data-dependent
 boolean mask

A boolean mask whose True count is only known at runtime makes the
number of updated rows a runtime value N in the index_put converter.
Values and indices with a static row count then failed or gave wrong
results:

- h[mask, :] += v on a 3-D h (mask over two dims), as transformers'
  Qwen3-VL adds its visual embeddings in _deepstack_process, failed the
  build with "Cannot broadcast (128, 2048) to (<ITensor>, 2048)". This
  breaks compiling GR00T, LaWAM and VLA-JEPA.
- h[mask] = v on a 2-D h compiled, but every row came from v's first
  row.
- h[mask, idx] = v with a static idx failed the same way as the first
  case, and once values were fixed, wrote every row at idx[0].

The cause is that impl.slice.expand broadcasts a static dim whose
target is a runtime value, which is only right for a dim of size 1.
PyTorch requires values and indices to have N rows there, so the
converter now keeps such a static first dim instead of expanding it to
N. Dims of size 1 still broadcast, other dims and static shapes are
unchanged. A mask whose True count differs from the number of value
rows now fails at runtime, as PyTorch does, instead of writing part of
the rows.
---
 .../dynamo/conversion/impl/select.py          | 46 +++++++++++--------
 .../dynamo/conversion/test_index_put_aten.py  | 39 ++++++++++++++++
 2 files changed, 66 insertions(+), 19 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py
index d473842cc81..37e2e45c71d 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/select.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py
@@ -683,6 +683,29 @@ def index_put_scatter_add_plugin(
     return out
 
 
+def _broadcast_target(
+    expected_shape: Sequence[Union[int, ITensor]], tensor_shape: Sequence[int]
+) -> List[Union[int, ITensor]]:
+    """The shape to expand a tensor laid out as (N, ...) to, broadcasting like PyTorch.
+
+    expand broadcasts any static dim whose target is an ITensor, which is only right
+    for size 1. A data-dependent N, such as the True count of a boolean mask, must
+    match a larger static first dim at runtime in PyTorch, so keep that dim.
+    """
+    tensor_shape = [1] * (len(expected_shape) - len(tensor_shape)) + list(tensor_shape)
+    shape: List[Union[int, ITensor]] = []
+    for k, (exp_dim, dim) in enumerate(zip(expected_shape, tensor_shape)):
+        if k == 0 and isinstance(exp_dim, ITensor) and dim not in (1, DYNAMIC_DIM):
+            shape.append(dim)
+        elif dim in (1, DYNAMIC_DIM) or exp_dim == dim:
+            shape.append(exp_dim)
+        else:
+            raise ValueError(
+                f"Cannot broadcast {tuple(tensor_shape)} to {expected_shape}"
+            )
+    return shape
+
+
 def index_put_converter(
     ctx: ConversionContext,
     target: Target,
@@ -853,7 +876,7 @@ def index_put_converter(
             source_ir,
             f"{name}_expand_idx_I_{i}",
             idx_reshaped,
-            (N, F_volume),
+            tuple(_broadcast_target((N, F_volume), idx_reshaped.shape)),
         )
         I_tensors.append(expanded_idx)
 
@@ -999,8 +1022,7 @@ def index_put_converter(
 
     # Define the expected shape based on (N,) + F_shape_values.
     # N may be a ITensor when it comes from a dynamic source (e.g. nonzero).
-    # Pass it directly — impl.slice.expand handles ITensor shape elements and
-    # will emit a stride-0 broadcast when expanding a size-1 dim to a dynamic N.
+    # _broadcast_target decides which values dims expand to it.
     expected_shape = (N,) + tuple(F_shape_values)
 
     # Broadcast 'values' to match the expected shape
@@ -1037,7 +1059,7 @@ def index_put_converter(
                 source_ir,
                 f"{name}_expand_values",
                 values_permuted,
-                expected_shape,
+                tuple(_broadcast_target(expected_shape, values_permuted.shape)),
             )
         elif K == 1 and len(values.shape) == 1 and F and I[0] > max(F):
             # For data[..., idx] = values, a 1-D values tensor carries the
@@ -1114,27 +1136,13 @@ def index_put_converter(
             )
         else:
             # Discontinuous case (K > 1 or K == 0)
-            values_shape_padded = [1] * (
-                len(expected_shape) - len(values.shape)
-            ) + list(values.shape)
-            broadcast_shape = []
-            for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
-                if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM:
-                    broadcast_shape.append(-1)
-                elif val_dim == 1 or exp_dim == val_dim:
-                    broadcast_shape.append(exp_dim)
-                else:
-                    raise ValueError(
-                        f"Cannot broadcast {values.shape} to {expected_shape}"
-                    )
-
             values_expanded = impl.slice.expand(
                 ctx,
                 target,
                 source_ir,
                 f"{name}_expand_values",
                 values,
-                expected_shape,
+                tuple(_broadcast_target(expected_shape, values.shape)),
             )
 
     # Flatten values to (N * F_volume,)
diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py
index 1460e5657c9..4ccf1747fbf 100644
--- a/tests/py/dynamo/conversion/test_index_put_aten.py
+++ b/tests/py/dynamo/conversion/test_index_put_aten.py
@@ -518,6 +518,45 @@ def forward(self, source_tensor, indices_tensor, value_tensor):
 
         torch.allclose(result, torch_output, atol=1e-4, rtol=1e-4)
 
+    def test_bool_mask_with_data_dependent_count(self):
+        # The True count is only known at runtime, while values has a static row count.
+        class AddRows(torch.nn.Module):
+            def forward(self, data, mask, values):
+                data = data.clone()
+                data[mask, :] += values  # Qwen3-VL adds visual embeddings this way
+                return data
+
+        class SetRows(torch.nn.Module):
+            def forward(self, data, mask, values):
+                data = data.clone()
+                data[mask] = values
+                return data
+
+        class SetRowsAtColumns(torch.nn.Module):
+            def forward(self, data, mask, values):
+                data = data.clone()
+                data[mask, torch.tensor([0, 2, 4, 5], device=data.device)] = values
+                return data
+
+        values = torch.randn(4, 8, device="cuda")
+        mask = torch.zeros(1, 10, dtype=torch.bool, device="cuda")
+        mask[0, 3:7] = True
+        cases = [
+            (AddRows(), torch.randn(1, 10, 8, device="cuda"), mask),
+            (SetRows(), torch.randn(10, 8, device="cuda"), mask[0]),
+            (SetRowsAtColumns(), torch.randn(10, 6, 8, device="cuda"), mask[0]),
+        ]
+        for model, data, case_mask in cases:
+            with self.subTest(type(model).__name__):
+                inputs = (data, case_mask, values)
+                compiled = torchtrt.dynamo.compile(
+                    torch.export.export(model, inputs),
+                    arg_inputs=inputs,
+                    min_block_size=1,
+                    pass_through_build_failures=True,
+                )
+                torch.testing.assert_close(compiled(*inputs), model(*inputs))
+
     def test_updates_are_cast_to_int64_destination_dtype(self):
         class IndexPutArgmax(torch.nn.Module):
             def forward(self, data, idx, scores):

From a35077bfd7e7459c689ef2d8a4e68be15aa65775 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 00:41:31 -0700
Subject: [PATCH 11/36] fix(executorch): export engines whose data-dependent
 shapes stay inside

Torch-TensorRT marks an engine as needing an output allocator as soon as
one of its operations has a data-dependent shape, such as aten.index.Tensor
with a boolean mask. Its runtimes use that mark to keep CUDA graphs off,
which is right: TensorRT cannot capture such an engine even when every
output shape is known. ExecuTorch export read the same mark as "an output
has a data-dependent shape" and refused the engine with "ExecuTorch export
does not support TensorRT engines that require an output allocator".

That is too broad. Qwen3-VL adds its visual embeddings with
h[mask, :] += v: the selected rows have a data-dependent count, but the
engine returns a tensor shaped like h. GR00T and LaWAM compiled to
TensorRT but could not be saved to ExecuTorch. The ExecuTorch backend
allocates static outputs and falls back to plain enqueue when a CUDA
graph capture fails, so it can run such an engine.

For a marked engine, export now asks the live engine for its output
shapes at the node's example input shapes (infer_outputs, the call the
Python runtime's fake kernel makes) and keeps refusing it only if an
output shape is still unknown. Otherwise the engine's ExecuTorch record
carries no mark, so the backend accepts it. Nothing is serialized for
the check; it creates the engine's execution context if no run has
created it yet. An engine whose input shapes have no example value, or
that cannot report its outputs, is still refused, and so is a marked
engine handed to the backend without this export step.
---
 py/torch_tensorrt/executorch/_export_utils.py | 44 +++++++++++++++++++
 tests/py/dynamo/executorch/test_export.py     | 38 ++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py
index f4730fae149..db4e7c0a111 100644
--- a/py/torch_tensorrt/executorch/_export_utils.py
+++ b/py/torch_tensorrt/executorch/_export_utils.py
@@ -12,6 +12,10 @@
 from torch._library.fake_class_registry import FakeScriptObject
 from torch._subclasses.fake_tensor import is_fake
 from torch.export.graph_signature import InputKind
+from torch.fx.experimental.symbolic_shapes import (
+    guarding_hint_or_throw,
+    has_guarding_hint,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -92,6 +96,29 @@ def get_engine_info_from_state(
     return list(state[0] if isinstance(state, tuple) else state)
 
 
+def _outputs_depend_on_data(engine_obj: Any, node: Any) -> bool:
+    """Whether the engine leaves an output shape unknown for this node's input shapes.
+
+    Torch-TensorRT marks an engine as needing an output allocator when any of its
+    operations has a data-dependent shape, which also keeps CUDA graphs off for it.
+    Only an output with such a shape needs the allocator ExecuTorch lacks.
+    """
+    input_shapes = []
+    for arg in node.args[0]:
+        val = arg.meta.get("val") if isinstance(arg, torch.fx.Node) else None
+        if not isinstance(val, torch.Tensor):
+            return True
+        if not all(has_guarding_hint(dim) for dim in val.shape):
+            return True
+        input_shapes.append([guarding_hint_or_throw(dim) for dim in val.shape])
+    try:
+        output_shapes = engine_obj.infer_outputs(input_shapes)
+    except RuntimeError:
+        # The C++ engine cannot create a tensor for an unknown dimension.
+        return True
+    return any(-1 in tuple(getattr(shape, "shape", shape)) for shape in output_shapes)
+
+
 def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> None:
     """Reject engine configurations unsupported by ExecuTorch export."""
     from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
@@ -237,6 +264,10 @@ def validate_engine_program(
     which never reads a cached record's ENGINE_IDX -- so a metadata-only record is
     never cross-served as the engine bytes.
     """
+    from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
+        REQUIRES_OUTPUT_ALLOCATOR_IDX,
+    )
+
     count = 0
     # One engine can feed several execute_engine nodes, and resolving it again would
     # serialize it again and keep a second copy of the bytes.
@@ -263,6 +294,19 @@ def validate_engine_program(
             )
             if engine_node is not None:
                 by_engine_node[engine_node] = engine_info
+        if (
+            node.target is torch.ops.tensorrt.execute_engine.default
+            and len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX
+            and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1"
+            and not _outputs_depend_on_data(
+                _resolve_engine_object(exported_program, node), node
+            )
+        ):
+            # Every output shape is known, so ExecuTorch needs no output allocator.
+            # This record becomes the ExecuTorch engine's arguments, and the backend
+            # refuses a marked one.
+            engine_info = list(engine_info)
+            engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "0"
         validate_engine_info(engine_info, node_name=node.name)
         if resolved is not None:
             resolved[node.name] = engine_info
diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py
index 8e72788e332..d95e987a586 100644
--- a/tests/py/dynamo/executorch/test_export.py
+++ b/tests/py/dynamo/executorch/test_export.py
@@ -857,6 +857,44 @@ def forward(self, value):
     assert output.stat().st_size > 0
 
 
+def test_engine_with_data_dependent_shape_only_inside_exports_a_pte(tmp_path):
+    """The boolean mask selection has a data-dependent shape, the engine output does not."""
+
+    class MaskedSum(torch.nn.Module):
+        def forward(self, value, mask):
+            return value + value[mask].sum()
+
+    class MaskedRows(torch.nn.Module):
+        def forward(self, value, mask):
+            return value[mask] * 2
+
+    value = torch.randn(10, 4, device="cuda")
+    mask = torch.arange(10, device="cuda") % 3 == 0
+
+    def save(model):
+        exported = torch.export.export(model.eval().cuda(), (value, mask))
+        compiled = torch_tensorrt.dynamo.compile(
+            exported,
+            arg_inputs=[value, mask],
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        output = tmp_path / f"{type(model).__name__}.pte"
+        torch_tensorrt.save(
+            compiled,
+            str(output),
+            output_format="executorch",
+            arg_inputs=(value, mask),
+            retrace=False,
+        )
+        return output
+
+    assert save(MaskedSum()).stat().st_size > 0
+    # An output shape that depends on the data still needs an output allocator.
+    with pytest.raises(RuntimeError, match="output allocator"):
+        save(MaskedRows())
+
+
 @pytest.mark.unit
 def test_lowering_stores_the_staged_engine_bytes_as_named_data(monkeypatch, caplog):
     """The rewrite stages each engine as ``bytes``, and lowering must hand that same

From b234948ad48a53ece3ecbb132a7a6114ded4a26e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 06:40:37 -0700
Subject: [PATCH 12/36] fix(conversion): keep select of a uint8 tensor in
 PyTorch

TensorRT allows uint8 tensors only at engine inputs and at outputs made
by a cast or identity layer. The select converter emits a gather, so
selecting from a uint8 input, such as frames[1] on a uint8 camera
history before casting it to float, made a uint8 intermediate and failed
the engine build:

  Tensor [GATHER]-[aten_ops.select.int]-[/select_gather]_output has
  datatype UInt8 but is not network-level IO

A select whose uint8 result is itself an engine output failed the same
way. VQ-BeT hits this with its uint8 image history. aten.gather already
keeps a uint8 input in PyTorch for the same reason, and select now does
too. Slice, permute, view and index on a uint8 input still fail; they
need the float cast moved ahead of them, which is a separate change.
---
 .../dynamo/conversion/aten_ops_converters.py  | 22 ++++++++++++++++++-
 .../py/dynamo/conversion/test_select_aten.py  | 17 ++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 4654518b90e..5fe7cd1ccdd 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1753,7 +1753,27 @@ def aten_ops_scatter(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.select.int, supports_dynamic_shapes=True)
+def select_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    """Keep a select of a uint8 tensor on the PyTorch path.
+
+    TensorRT allows uint8 only at engine inputs and at outputs made by a cast or
+    identity layer, so the gather's uint8 result fails the engine build, as it does
+    for aten.gather.
+    """
+    meta = getattr(args_bounds_check(node.args, 0), "meta", {}).get("val")
+    if isinstance(meta, torch.Tensor) and meta.dtype == torch.uint8:
+        _LOGGER.debug("select of a uint8 tensor is not supported, falling back")
+        return False
+    return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.select.int,
+    capability_validator=select_validator,
+    supports_dynamic_shapes=True,
+)
 def aten_ops_select(
     ctx: ConversionContext,
     target: Target,
diff --git a/tests/py/dynamo/conversion/test_select_aten.py b/tests/py/dynamo/conversion/test_select_aten.py
index 15eb1d82a69..952f63c108e 100644
--- a/tests/py/dynamo/conversion/test_select_aten.py
+++ b/tests/py/dynamo/conversion/test_select_aten.py
@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
@@ -163,6 +164,22 @@ def forward(self, values, index):
                 check_dtype=False,
             )
 
+    def test_select_uint8_input_compiles(self):
+        # The uint8 select would fail the engine build, so it must stay in PyTorch.
+        class SelectThenCast(nn.Module):
+            def forward(self, frames):
+                return frames[1].float() / 255
+
+        frames = torch.randint(0, 256, (2, 8, 8, 3), dtype=torch.uint8, device="cuda")
+        model = SelectThenCast()
+        compiled = torch_tensorrt.dynamo.compile(
+            torch.export.export(model, (frames,)),
+            arg_inputs=[frames],
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        torch.testing.assert_close(compiled(frames), model(frames))
+
 
 if __name__ == "__main__":
     run_tests()

From 18b4140f2aaefa092f39954ee2c3d1d1a99e212a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 09:09:21 -0700
Subject: [PATCH 13/36] fix(conversion): read tensor_meta in the uint8 select
 check

select_validator read only node.meta["val"]. A graph typed by torch.fx
ShapeProp instead of torch.export, for example one built with
symbolic_trace, carries tensor_meta but no val, so the check missed a
uint8 input there. It now falls back to tensor_meta, as gather_validator
already does. This builds on the previous commit, which added the check.
---
 .../dynamo/conversion/aten_ops_converters.py        |  4 ++--
 tests/py/dynamo/conversion/test_select_aten.py      | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 5fe7cd1ccdd..49d0fde4503 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1762,8 +1762,8 @@ def select_validator(
     identity layer, so the gather's uint8 result fails the engine build, as it does
     for aten.gather.
     """
-    meta = getattr(args_bounds_check(node.args, 0), "meta", {}).get("val")
-    if isinstance(meta, torch.Tensor) and meta.dtype == torch.uint8:
+    meta = getattr(args_bounds_check(node.args, 0), "meta", {})
+    if getattr(meta.get("val", meta.get("tensor_meta")), "dtype", None) == torch.uint8:
         _LOGGER.debug("select of a uint8 tensor is not supported, falling back")
         return False
     return True
diff --git a/tests/py/dynamo/conversion/test_select_aten.py b/tests/py/dynamo/conversion/test_select_aten.py
index 952f63c108e..260e6f6ab8d 100644
--- a/tests/py/dynamo/conversion/test_select_aten.py
+++ b/tests/py/dynamo/conversion/test_select_aten.py
@@ -5,8 +5,10 @@
 import torch.nn as nn
 import torch_tensorrt
 from parameterized import parameterized
+from torch.fx.passes.shape_prop import ShapeProp
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.conversion.aten_ops_converters import select_validator
 
 from .harness import DispatchTestCase
 
@@ -180,6 +182,17 @@ def forward(self, frames):
         )
         torch.testing.assert_close(compiled(frames), model(frames))
 
+    def test_select_validator_reads_tensor_meta(self):
+        class Select(nn.Module):
+            def forward(self, frames):
+                return torch.ops.aten.select.int(frames, 0, 1)
+
+        graph_module = torch.fx.symbolic_trace(Select())
+        ShapeProp(graph_module).propagate(torch.zeros(2, 4, dtype=torch.uint8))
+        node = next(n for n in graph_module.graph.nodes if n.op == "call_function")
+        self.assertNotIn("val", node.args[0].meta)
+        self.assertFalse(select_validator(node))
+
 
 if __name__ == "__main__":
     run_tests()

From 47b40be5a9ea1088d6b4401da0a2d3d43520776c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 14:11:35 -0700
Subject: [PATCH 14/36] fix(lowering): run complex prod in PyTorch on the
 legacy complex rewriter

The legacy complex rewriter stores a complex tensor as a real tensor with a
trailing dimension of 2. For `prod.dim_int` it kept the op on that layout, so
it multiplied the real parts together and the imaginary parts together. That
is not a complex product, so a product of two or more values with imaginary
parts came out wrong.

There is no exact way to write a complex product over a dimension of any
length in real arithmetic. Going through magnitudes and angles loses
precision: a product of 4096 values of -1 drifts away from 1, and
1e-30 * 1e15 * 1e15 underflows to 0 when the parts are squared. So the
rewriter no longer has a rule for `prod.dim_int`. Like other complex ops
without a rule, it now runs in PyTorch on genuine complex tensors, and the
ops around it still go to TensorRT.

Test: a prod over a positive dim, and over a negative dim with keepdim, static
and with a dynamic batch, compared with eager at two batch sizes.
---
 .../lowering/passes/complex_graph_rewrite.py  | 17 --------
 tests/py/dynamo/hlo/test_complex_ops.py       | 39 +++++++++++++++++++
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
index be9134e33d4..77cc189b1ee 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
@@ -711,23 +711,6 @@ def _rewrite_mean_dim(self, node: Node) -> bool:
             self.gm.graph.erase_node(node)
             return True
 
-    @_complex_unpacker(torch.ops.aten.prod.dim_int)
-    def _rewrite_prod_dim(self, node: Node) -> bool:
-        # prod.dim_int(inp, dim, keepdim=False, dtype=None)
-        inp = node.args[0]
-        dim = node.args[1]
-        if dim >= 0:
-            return False  # positive dim — pass-through is correct
-        new_dim = dim - 1
-        keepdim = node.args[2] if len(node.args) > 2 else False
-        extra = list(node.args[3:])
-        with SubgraphBuilder(self.gm.graph, node) as b:
-            out = b(torch.ops.aten.prod.dim_int, inp, new_dim, keepdim, *extra)
-            node.replace_all_uses_with(out)
-            out.meta["is_complex_layout"] = True
-            self.gm.graph.erase_node(node)
-            return True
-
     @_complex_unpacker(torch.ops.aten.narrow.default)
     def _rewrite_narrow(self, node: Node) -> bool:
         # narrow(inp, dim, start, length) — shift negative dim by -1
diff --git a/tests/py/dynamo/hlo/test_complex_ops.py b/tests/py/dynamo/hlo/test_complex_ops.py
index 75e1f04262c..6b79bbb9e58 100644
--- a/tests/py/dynamo/hlo/test_complex_ops.py
+++ b/tests/py/dynamo/hlo/test_complex_ops.py
@@ -905,6 +905,45 @@ def test_complex_output_dynamic_batch():
     torch._dynamo.reset()
 
 
+class ComplexOp(nn.Module):
+    """Apply one complex op to a real (..., 2) input viewed as complex."""
+
+    def __init__(self, op):
+        super().__init__()
+        self.op = op
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.view_as_real(self.op(torch.view_as_complex(x)))
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("dynamic", [False, True], ids=["static", "dynamic"])
+@pytest.mark.parametrize(
+    "op",
+    [
+        lambda z: torch.prod(z, dim=1),
+        lambda z: torch.prod(z, dim=-1, keepdim=True),
+    ],
+    ids=["prod_dim", "prod_neg_dim_keepdim"],
+)
+def test_complex_op_matches_eager(op, dynamic):
+    model = ComplexOp(op).eval().cuda()
+    x = torch.randn(2, 4, 2).cuda()
+    dynamic_shapes = ({0: Dim("batch", min=2, max=8)},) if dynamic else None
+    ep = torch.export.export(model, (x,), dynamic_shapes=dynamic_shapes)
+    trt_model = torchtrt.dynamo.compile(
+        ep,
+        inputs=[x],
+        min_block_size=1,
+        pass_through_build_failures=True,
+        **_DECOMP,
+    )
+    for batch in (2, 5) if dynamic else (2,):
+        xb = torch.randn(batch, 4, 2).cuda()
+        torch.testing.assert_close(trt_model(xb), model(xb), rtol=1e-3, atol=1e-3)
+    torch._dynamo.reset()
+
+
 # ===========================================================================
 # 11. Numerical precision: complex64 vs truncated complex128
 # ===========================================================================

From 93508ae1576187aa6621f2e5a90ffca63f883443 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 14:45:30 -0700
Subject: [PATCH 15/36] fix(lowering): handle complex Python scalars in the
 legacy complex rewriter

`z + (1+2j)`, `z * (2-1j)` and `z / (2-1j)` gave wrong values on the legacy
complex rewriter. It treated every Python scalar as real: a scalar added to the
real part only, and a scalar multiplied both parts equally. For a complex
scalar both are wrong. `z * (2-1j)` and `z / (2-1j)` also built no TensorRT
engine, because the complex scalar was left in the graph.

Now a complex scalar is split into its real and imaginary parts:

- add and sub apply each part to the matching part of the tensor;
- mul does a full complex multiply;
- div multiplies by the scalar's reciprocal, computed in Python; for 0j that
  is nan, as on CUDA eager, instead of a Python error.

The same holds when the tensor is real or an integer, as in `x * 1j`: its
imaginary part is zero. The `mul.Scalar` and `div.Scalar` overloads now go
through the same rule instead of passing through as if the scalar were real.
A symbolic scalar, such as a dynamic size under `torch.compile`, is always
real and still scales both parts.
Real scalars take the same path as before, except that add and sub now apply
`alpha`, which was dropped for every scalar. The `add.Scalar` and `sub.Scalar`
handler did the same thing as the `add.Tensor` scalar path, so the two are now
one handler.

Tests: add (also with alpha), sub, mul and div with a complex scalar, on a
complex, a real and an integer tensor, static and with a dynamic batch,
compared with eager; and a complex tensor divided by its dynamic batch size
under `torch.compile`.
---
 .../lowering/passes/complex_graph_rewrite.py  | 108 +++++++++++++-----
 tests/py/dynamo/hlo/test_complex_ops.py       |  47 +++++++-
 2 files changed, 124 insertions(+), 31 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
index 77cc189b1ee..5238a37526a 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
@@ -33,8 +33,6 @@
         torch.ops.aten.add.Tensor,
         torch.ops.aten.sub.Tensor,
         torch.ops.aten.neg.default,
-        torch.ops.aten.mul.Scalar,  # scalar*(re,im) — both parts scaled equally
-        torch.ops.aten.div.Scalar,  # (re,im)/scalar — both parts divided equally
         # Structural / copy — operate on the whole tensor without touching content.
         # Note: permute.default is NOT here; it needs an explicit rewrite to append
         # the trailing real/imag dimension index to the dims list.
@@ -463,7 +461,12 @@ def _rewrite_permute(self, node: Node) -> bool:
             self.gm.graph.erase_node(node)
             return True
 
-    @_complex_unpacker(torch.ops.aten.mul.Tensor, torch.ops.aten.div.Tensor)
+    @_complex_unpacker(
+        torch.ops.aten.mul.Tensor,
+        torch.ops.aten.div.Tensor,
+        torch.ops.aten.mul.Scalar,
+        torch.ops.aten.div.Scalar,
+    )
     def _rewrite_mul_div_tensor(self, node: Node) -> bool:
         arg0_is_node = isinstance(node.args[0], torch.fx.Node)
         arg1_is_node = isinstance(node.args[1], torch.fx.Node)
@@ -471,13 +474,29 @@ def _rewrite_mul_div_tensor(self, node: Node) -> bool:
         if not arg0_is_node and not arg1_is_node:
             return False  # both scalars
 
-        if node.target == torch.ops.aten.mul.Tensor and (
-            not arg0_is_node or not arg1_is_node
+        if node.target in (torch.ops.aten.mul.Scalar, torch.ops.aten.div.Scalar) and (
+            arg1_is_node
         ):
-            return False  # scalar * complex — elementwise-safe
-
-        if node.target == torch.ops.aten.div.Tensor and not arg1_is_node:
-            return False  # complex / scalar — elementwise-safe
+            return False  # symbolic real scalar: elementwise-safe
+
+        is_mul = node.target in (torch.ops.aten.mul.Tensor, torch.ops.aten.mul.Scalar)
+        if is_mul and (not arg0_is_node or not arg1_is_node):
+            z_node, scalar = (
+                (node.args[0], node.args[1])
+                if arg0_is_node
+                else (node.args[1], node.args[0])
+            )
+            if isinstance(scalar, complex):
+                return self._rewrite_mul_complex_scalar(node, z_node, scalar)
+            return False  # real scalar * complex: elementwise-safe
+
+        if not is_mul and not arg1_is_node:
+            c = node.args[1]
+            if isinstance(c, complex):
+                # Python raises on 1 / 0j; nan is what CUDA eager returns for it.
+                inverse = 1 / c if c else complex("nan+nanj")
+                return self._rewrite_mul_complex_scalar(node, node.args[0], inverse)
+            return False  # complex / real scalar: elementwise-safe
 
         if node.target == torch.ops.aten.div.Tensor and not arg0_is_node:
             # scalar / complex: s/(a+bi) = (s*a/(a²+b²)) + i*(-s*b/(a²+b²))
@@ -625,21 +644,63 @@ def match_complex_mul(
                 return True
         return False
 
-    @_complex_unpacker(torch.ops.aten.add.Tensor, torch.ops.aten.sub.Tensor)
-    def _rewrite_add_sub_tensor_scalar(self, node: Node) -> bool:
+    def _rewrite_mul_complex_scalar(self, node: Node, z_node: Node, c: complex) -> bool:
+        with SubgraphBuilder(self.gm.graph, node) as b:
+            if self._is_complex_layout_node(z_node):
+                # (a+bi)(x+yi) = (ax-by) + (ay+bx)i
+                re, im = self._inline_select_re_im(b, z_node)
+                re_x = b(torch.ops.aten.mul.Tensor, re, c.real)
+                im_y = b(torch.ops.aten.mul.Tensor, im, c.imag)
+                out_re = b(torch.ops.aten.sub.Tensor, re_x, im_y)
+                re_y = b(torch.ops.aten.mul.Tensor, re, c.imag)
+                im_x = b(torch.ops.aten.mul.Tensor, im, c.real)
+                out_im = b(torch.ops.aten.add.Tensor, re_y, im_x)
+            else:
+                # a(x+yi) = ax + ayi for a real tensor a
+                out_re = b(torch.ops.aten.mul.Tensor, z_node, c.real)
+                out_im = b(torch.ops.aten.mul.Tensor, z_node, c.imag)
+            out = self._inline_cat_re_im(b, out_re, out_im)
+            node.replace_all_uses_with(out)
+            self.gm.graph.erase_node(node)
+            return True
+
+    @_complex_unpacker(
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.sub.Tensor,
+        torch.ops.aten.add.Scalar,
+        torch.ops.aten.sub.Scalar,
+    )
+    def _rewrite_add_sub_scalar(self, node: Node) -> bool:
         # add.Tensor(z1, z2): tensor+tensor is componentwise on the [..., 2]
         # real layout — real and imag parts are handled independently — so no
         # rewrite is needed.  Return True (handled, no modification).
-        # add.Tensor(z, scalar): scalar applies to real part only, so we must
-        # explicitly rewrite.
-        # tensor+tensor: componentwise on [..., 2] layout — no rewrite needed
+        # (a+bi) ± (x+yi) = (a±x) + (b±y)i, so a real scalar touches only the
+        # real part and a complex scalar adds its parts separately.
         if len(node.args) < 2 or isinstance(node.args[1], torch.fx.Node):
             return True
         inp, scalar = node.args[0], node.args[1]
+        scalar = scalar * node.kwargs.get("alpha", 1)
+        s_re, s_im = (
+            (scalar.real, scalar.imag) if isinstance(scalar, complex) else (scalar, 0)
+        )
         with SubgraphBuilder(self.gm.graph, node) as b:
-            re = b(torch.ops.aten.select.int, inp, -1, 0)
-            im = b(torch.ops.aten.select.int, inp, -1, 1)
-            new_re = b(node.target, re, scalar)
+            if self._is_complex_layout_node(inp):
+                re, im = self._inline_select_re_im(b, inp)
+                if s_im:
+                    im = b(node.target, im, s_im)
+            else:
+                re, im = inp, None
+            new_re = b(node.target, re, s_re)
+            if im is None:
+                # A real tensor has no imaginary part, so the result's is ±y. Fill
+                # it like new_re, which already has the promoted float dtype.
+                is_add = node.target in (
+                    torch.ops.aten.add.Tensor,
+                    torch.ops.aten.add.Scalar,
+                )
+                im = b(
+                    torch.ops.aten.full_like.default, new_re, s_im if is_add else -s_im
+                )
             out = self._inline_cat_re_im(b, new_re, im)
             node.replace_all_uses_with(out)
             out.meta["is_complex_layout"] = True
@@ -912,19 +973,6 @@ def _rewrite_polar(self, node: Node) -> bool:
             self.gm.graph.erase_node(node)
             return True
 
-    @_complex_unpacker(torch.ops.aten.add.Scalar, torch.ops.aten.sub.Scalar)
-    def _rewrite_add_sub_scalar(self, node: Node) -> bool:
-        # (a+bi) ± s = (a±s) + bi — scalar applies to real part only
-        inp, scalar = node.args[0], node.args[1]
-        with SubgraphBuilder(self.gm.graph, node) as b:
-            re = b(torch.ops.aten.select.int, inp, -1, 0)
-            im = b(torch.ops.aten.select.int, inp, -1, 1)
-            new_re = b(node.target, re, scalar)
-            out = self._inline_cat_re_im(b, new_re, im)
-            node.replace_all_uses_with(out)
-            self.gm.graph.erase_node(node)
-            return True
-
     @_complex_unpacker(torch.ops.aten.log2.default, torch.ops.aten.log10.default)
     def _rewrite_log2_log10(self, node: Node) -> bool:
         # log_b(z) = log(z) / log(b)
diff --git a/tests/py/dynamo/hlo/test_complex_ops.py b/tests/py/dynamo/hlo/test_complex_ops.py
index 6b79bbb9e58..31016018d6d 100644
--- a/tests/py/dynamo/hlo/test_complex_ops.py
+++ b/tests/py/dynamo/hlo/test_complex_ops.py
@@ -923,8 +923,29 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     [
         lambda z: torch.prod(z, dim=1),
         lambda z: torch.prod(z, dim=-1, keepdim=True),
+        lambda z: z + (1 + 2j),
+        lambda z: torch.ops.aten.sub.Scalar(z, 1 + 2j),
+        lambda z: (2 - 1j) * z,
+        lambda z: z / (2 - 1j),
+        lambda z: torch.ops.aten.mul.Scalar(z, 2 - 1j),
+        lambda z: z.real * (2 - 1j),
+        lambda z: z.imag - (1 + 2j),
+        lambda z: (z.real * 4).to(torch.int64) + (1.25 + 2.75j),
+        lambda z: torch.add(z, 1 + 2j, alpha=2),
+    ],
+    ids=[
+        "prod_dim",
+        "prod_neg_dim_keepdim",
+        "add_complex_scalar",
+        "sub_complex_scalar",
+        "mul_complex_scalar",
+        "div_complex_scalar",
+        "mul_scalar_overload_complex_scalar",
+        "real_tensor_mul_complex_scalar",
+        "real_tensor_sub_complex_scalar",
+        "int_tensor_add_complex_scalar",
+        "add_complex_scalar_alpha",
     ],
-    ids=["prod_dim", "prod_neg_dim_keepdim"],
 )
 def test_complex_op_matches_eager(op, dynamic):
     model = ComplexOp(op).eval().cuda()
@@ -944,6 +965,30 @@ def test_complex_op_matches_eager(op, dynamic):
     torch._dynamo.reset()
 
 
+class ComplexDivBySize(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        z = torch.view_as_complex(x)
+        size = torch.sym_float(x.shape[0])
+        return torch.view_as_real(torch.ops.aten.div.Scalar(z, size))
+
+
+@pytest.mark.unit
+def test_complex_divided_by_symbolic_size():
+    # Under torch.compile with dynamic shapes the divisor stays a graph value.
+    model = ComplexDivBySize().eval().cuda()
+    compiled = torch.compile(
+        model,
+        backend="torch_tensorrt",
+        fullgraph=True,
+        dynamic=True,
+        options={"min_block_size": 1, "pass_through_build_failures": True, **_DECOMP},
+    )
+    for batch in (3, 5):
+        x = torch.randn(batch, 4, 2).cuda()
+        torch.testing.assert_close(compiled(x), model(x), rtol=1e-3, atol=1e-3)
+    torch._dynamo.reset()
+
+
 # ===========================================================================
 # 11. Numerical precision: complex64 vs truncated complex128
 # ===========================================================================

From df50e35cd61155b137a1187bedd2bdb8fecd8e8f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 12:50:59 -0700
Subject: [PATCH 16/36] fix(lowering): rewrite torch.complex in the legacy
 complex rewriter

The legacy complex rewriter had no rule for `aten.complex`, which builds a
complex tensor from a real part and an imaginary part. The op stayed in the
graph as a complex tensor while the ops after it expected the real layout with
a trailing dimension of 2, so `torch.complex(z.real * 2, z.imag)` returned
shape `(b, 4)` instead of `(b, 4, 2)`.

The new rule stacks the two parts on a new last dimension, the same way the
complex decomposition path packs `aten.complex`.

Test: `torch.complex` of a scaled real part and the imaginary part, static and
with a dynamic batch, compared with eager.
---
 .../dynamo/lowering/passes/complex_graph_rewrite.py      | 9 +++++++++
 tests/py/dynamo/hlo/test_complex_ops.py                  | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
index 5238a37526a..dd365cfcb4e 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py
@@ -959,6 +959,15 @@ def _rewrite_angle(self, node: Node) -> bool:
             self.gm.graph.erase_node(node)
             return True
 
+    @_complex_unpacker(torch.ops.aten.complex.default)
+    def _rewrite_complex(self, node: Node) -> bool:
+        # complex(re, im) is the [..., 2] layout of its two real inputs.
+        with SubgraphBuilder(self.gm.graph, node) as b:
+            out = self._inline_cat_re_im(b, node.args[0], node.args[1])
+            node.replace_all_uses_with(out)
+            self.gm.graph.erase_node(node)
+            return True
+
     @_complex_unpacker(torch.ops.aten.polar.default)
     def _rewrite_polar(self, node: Node) -> bool:
         # polar(r, theta) = r*cos(theta) + i*r*sin(theta)
diff --git a/tests/py/dynamo/hlo/test_complex_ops.py b/tests/py/dynamo/hlo/test_complex_ops.py
index 31016018d6d..9f2351afb26 100644
--- a/tests/py/dynamo/hlo/test_complex_ops.py
+++ b/tests/py/dynamo/hlo/test_complex_ops.py
@@ -932,6 +932,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         lambda z: z.imag - (1 + 2j),
         lambda z: (z.real * 4).to(torch.int64) + (1.25 + 2.75j),
         lambda z: torch.add(z, 1 + 2j, alpha=2),
+        lambda z: torch.complex(z.real * 2, z.imag),
     ],
     ids=[
         "prod_dim",
@@ -945,6 +946,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         "real_tensor_sub_complex_scalar",
         "int_tensor_add_complex_scalar",
         "add_complex_scalar_alpha",
+        "complex_from_parts",
     ],
 )
 def test_complex_op_matches_eager(op, dynamic):

From 4705c23a4711c1faf859960795b373721bae2487 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 13:01:53 -0700
Subject: [PATCH 17/36] fix(refit): give the output check's reference module
 complex inputs and outputs

`refit_module_weights(..., verify_output=True)` crashed for a model with a
complex input when it was compiled with `use_complex_decomposition=False`:

    RuntimeError: The size of tensor a (2) must match the size of tensor b (4)

And for a model with a complex output it logged "Refitting Failed! The
outputs do not match." on both complex paths, although the refit was correct.

The check runs the caller's inputs through the lowered PyTorch module as a
reference. That module does not have the compiled module's complex boundary:
the legacy complex rewriter replaces a complex input with a real
`<name>_unpacked_complex` one, and both complex lowerings return a complex
output as a real tensor with a trailing dimension of 2.

The check now adds that boundary to the reference module, from the input
names and output indices the complex lowering records: `view_as_real` after
each unpacked input and `view_as_complex` on each complex output, as at the
compiled module's boundary.

Test: the complex buffer refit test now runs on both complex paths, static and
with a dynamic batch, with a real and a complex output, and checks that the
output check reports success.
---
 py/torch_tensorrt/dynamo/_refit.py         | 40 ++++++++++++++++++++--
 tests/py/dynamo/models/test_model_refit.py | 18 ++++++++--
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/_refit.py b/py/torch_tensorrt/dynamo/_refit.py
index 3762b3aed0c..79b3ab4b1bc 100644
--- a/py/torch_tensorrt/dynamo/_refit.py
+++ b/py/torch_tensorrt/dynamo/_refit.py
@@ -460,9 +460,7 @@ def refit_module_weights(
 
     if verify_output and arg_inputs is not None:
         new_gm.to(to_torch_device(settings.device))  # move to device for inference
-        # A complex placeholder stays complex in the lowered reference module; the
-        # unpacking to real happens inside the graph, so both modules take the same
-        # inputs the caller passed.
+        _restore_complex_io(new_gm)
         outputs_match = check_module_output(
             new_module=new_gm,
             refitted_module=compiled_module,
@@ -479,6 +477,42 @@ def refit_module_weights(
     return compiled_module
 
 
+def _restore_complex_io(gm: torch.fx.GraphModule) -> None:
+    """Make the lowered reference module take and return complex tensors, as the
+    compiled module does, so the output check can compare the two.
+
+    The legacy complex rewriter replaces a complex input with a real
+    "<name>_unpacked_complex" one, and both complex lowerings return a complex
+    output as a real tensor with a trailing dimension of 2.
+    """
+    graph = gm.graph
+    unpacked = {f"{n}_unpacked_complex" for n in gm.meta.get("complex_input_names", [])}
+    for node in list(graph.nodes):
+        if node.op == "placeholder" and node.name in unpacked:
+            with graph.inserting_after(node):
+                real = graph.call_function(torch.ops.aten.view_as_real.default, (node,))
+            node.replace_all_uses_with(real)
+            real.args = (node,)
+    output_node = list(graph.nodes)[-1]
+    output_arg = output_node.args[0]
+    outputs = (
+        list(output_arg) if isinstance(output_arg, (list, tuple)) else [output_arg]
+    )
+    for idx in gm.meta.get("complex_output_indices", []):
+        with graph.inserting_before(output_node):
+            outputs[idx] = graph.call_function(
+                torch.ops.aten.view_as_complex.default, (outputs[idx],)
+            )
+    output_node.args = (
+        (
+            type(output_arg)(outputs)
+            if isinstance(output_arg, (list, tuple))
+            else outputs[0]
+        ),
+    )
+    gm.recompile()
+
+
 # Util functions -----------
 import base64
 
diff --git a/tests/py/dynamo/models/test_model_refit.py b/tests/py/dynamo/models/test_model_refit.py
index 491283088fb..a6b1b19eb64 100644
--- a/tests/py/dynamo/models/test_model_refit.py
+++ b/tests/py/dynamo/models/test_model_refit.py
@@ -14,6 +14,7 @@
 import torch_tensorrt as torchtrt
 import torch_tensorrt as torch_trt
 from torch import nn
+from torch_tensorrt._features import has_complex_decomposition
 from torch_tensorrt._utils import is_tensorrt_rtx_version_supported
 from torch_tensorrt.dynamo import partitioning, refit_module_weights
 from torch_tensorrt.dynamo._refit import (
@@ -1037,8 +1038,16 @@ def forward(self, x):
     "Refit feature is not supported in Python 3.13 or higher",
 )
 @pytest.mark.unit
+@pytest.mark.parametrize(
+    "complex_output", [False, True], ids=["real_out", "complex_out"]
+)
 @pytest.mark.parametrize("dynamic_batch", [False, True])
-def test_complex_buffer_refit(dynamic_batch, caplog):
+@pytest.mark.parametrize(
+    "use_complex_decomposition", [False, True], ids=["legacy", "decomp"]
+)
+def test_complex_buffer_refit(
+    use_complex_decomposition, dynamic_batch, complex_output, caplog
+):
     """Refit a model whose weights include a complex-valued buffer (e.g. RoPE freqs).
 
     Exercises the combined complex_graph_detection + refit_module_weights path:
@@ -1046,6 +1055,8 @@ def test_complex_buffer_refit(dynamic_batch, caplog):
       - complex placeholder input goes through view_as_real at the TRT boundary
       - after refitting with new frequencies the output matches the new model
     """
+    if use_complex_decomposition and not has_complex_decomposition():
+        pytest.skip("decompose_complex_in_graph requires torch>=2.14.dev")
 
     class ComplexFreqModel(nn.Module):
         def __init__(self, freqs: torch.Tensor):
@@ -1053,8 +1064,8 @@ def __init__(self, freqs: torch.Tensor):
             self.register_buffer("freqs", freqs.cuda())
 
         def forward(self, z: torch.Tensor) -> torch.Tensor:
-            # complex mul then expose as real so TRT can produce a real output
-            return torch.view_as_real(z * self.freqs)
+            out = z * self.freqs
+            return out if complex_output else torch.view_as_real(out)
 
     SEQ, DIM = 8, 32
 
@@ -1087,6 +1098,7 @@ def make_freqs() -> torch.Tensor:
         tuple(inputs),
         min_block_size=1,
         immutable_weights=False,
+        use_complex_decomposition=use_complex_decomposition,
     )
 
     with caplog.at_level(logging.INFO, logger="torch_tensorrt.dynamo._refit"):

From b618b182df5946a23ce4701ec5481127002da23d Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 14:55:18 -0700
Subject: [PATCH 18/36] fix(executorch): fail cleanly when the TensorRT
 delegate has no GPU

Loading a .pte that holds a TensorRT engine on a machine with no visible
GPU, for example under CUDA_VISIBLE_DEVICES="", crashed the process with
a segmentation fault. Before loading, ExecuTorch asks the delegate
whether it is available, and the delegate answers by creating its shared
TensorRT runtime. Without a usable CUDA device, createInferRuntime logs
"CUDA initialization failure" and then returns a pointer that is not
null but is not a runtime either. The delegate then crashed calling it.

The delegate now asks CUDA for the device count before it creates the
runtime. When CUDA fails, it logs the reason and reports itself as not
available, so the load returns Error::NotFound with this in the log:

  TensorRTBackend: cannot create the TensorRT runtime, because CUDA is
  unavailable: no CUDA-capable device is detected

The new test exports a small program, loads it in a subprocess with no
visible GPU, and checks for this error. Without the fix the subprocess
dies with a segmentation fault.
---
 cpp/src/torch_tensorrt/executorch/README.md   |  3 +
 .../executorch/TensorRTBackend.cpp            | 12 +++
 .../py/dynamo/executorch/test_load_errors.py  | 79 +++++++++++++++++++
 3 files changed, 94 insertions(+)
 create mode 100644 tests/py/dynamo/executorch/test_load_errors.py

diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md
index f607cc7f1e8..125f28a73d8 100644
--- a/cpp/src/torch_tensorrt/executorch/README.md
+++ b/cpp/src/torch_tensorrt/executorch/README.md
@@ -194,6 +194,9 @@ but corrupted engine bytes can load and return wrong answers.
 
 ### Load errors
 
+- With no usable CUDA device, for example under `CUDA_VISIBLE_DEVICES=""`, the delegate reports
+  itself unavailable, so loading returns `Error::NotFound`. The log says CUDA is unavailable and
+  gives CUDA's reason.
 - Loading returns `Error::MemoryAllocationFailed` when TensorRT runs out of memory while it
   deserializes the engine, applies the weight streaming budget, or creates the execution context.
   The log then shows the free and total GPU memory. Other TensorRT load failures return
diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
index dca7601b46e..c5e49f3679b 100644
--- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
+++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
@@ -253,6 +253,18 @@ nvinfer1::IRuntime* TensorRTBackend::shared_runtime() {
   // Retried while null, so one transient failure does not disable the backend for the process.
   if (runtime == nullptr) {
     const std::shared_lock<std::shared_mutex> capture_lock(cuda_graph_capture_mutex());
+    // Without a usable CUDA device, createInferRuntime logs an error and then returns a pointer that
+    // is not null and is not a runtime either (seen with TensorRT 11.4), so check CUDA first.
+    int device_count = 0;
+    const cudaError_t cuda_err = cudaGetDeviceCount(&device_count);
+    if (cuda_err != cudaSuccess) {
+      cudaGetLastError();
+      ET_LOG(
+          Error,
+          "TensorRTBackend: cannot create the TensorRT runtime, because CUDA is unavailable: %s",
+          cudaGetErrorString(cuda_err));
+      return nullptr;
+    }
     if (logger == nullptr) {
       logger = new TRTLogger();
     }
diff --git a/tests/py/dynamo/executorch/test_load_errors.py b/tests/py/dynamo/executorch/test_load_errors.py
new file mode 100644
index 00000000000..8ae74939470
--- /dev/null
+++ b/tests/py/dynamo/executorch/test_load_errors.py
@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Load real TensorRT programs in ways that have to fail, and check each fails with a clear error."""
+
+import importlib.util
+import os
+import subprocess
+import sys
+
+import pytest
+
+pytest.importorskip("executorch.exir")
+
+import torch  # noqa: E402
+
+
+@pytest.fixture(autouse=True)
+def requires_runtime():
+    # Decided at run time, not collection time, for the reason test_cuda_partitioner_composition.py
+    # gives: remote GPU runners collect off the GPU host.
+    if not torch.cuda.is_available():
+        pytest.skip("needs a CUDA device to build a TensorRT engine")
+    if importlib.util.find_spec("torch_tensorrt_executorch_runtime") is None:
+        pytest.skip(
+            "needs the torch-tensorrt-executorch-runtime package to load the program"
+        )
+
+
+def _save(path, **kwargs):
+    import torch_tensorrt
+
+    torch.manual_seed(0)
+    inputs = (torch.randn(4, 16, device="cuda"),)
+    model = torch.nn.Linear(16, 16).eval().cuda()
+    trt_gm = torch_tensorrt.dynamo.compile(
+        torch.export.export(model, inputs), inputs=list(inputs), min_block_size=1
+    )
+    torch_tensorrt.save(
+        trt_gm,
+        str(path),
+        output_format="executorch",
+        retrace=False,
+        arg_inputs=list(inputs),
+        **kwargs,
+    )
+
+
+def test_loading_with_no_visible_gpu_fails_with_a_clear_error(tmp_path):
+    from executorch.exir import ExecutorchBackendConfig
+
+    pte = tmp_path / "model.pte"
+    # Planned in host memory, so loading reaches the delegate. A device-planned buffer would stop
+    # the load earlier, in ExecuTorch's CUDA allocator, which reports a missing GPU on its own.
+    _save(
+        pte,
+        backend_config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=False),
+    )
+    script = (
+        "import sys\n"
+        "from pathlib import Path\n"
+        "import torch_tensorrt_executorch_runtime\n"
+        "from executorch.runtime import Runtime\n"
+        "program = Runtime.get().load_program(Path(sys.argv[1]))\n"
+        "try:\n"
+        "    program.load_method('forward')\n"
+        "except RuntimeError as error:\n"
+        "    print(error)\n"
+        "    sys.exit(3)\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", script, str(pte)],
+        env={**os.environ, "CUDA_VISIBLE_DEVICES": ""},
+        capture_output=True,
+        text=True,
+        timeout=600,
+    )
+    assert result.returncode == 3, (result.returncode, result.stderr[-2000:])
+    assert "CUDA is unavailable" in result.stdout + result.stderr

From 17faf1e5f225da2db3717900e443dafe339c6a7d Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:08:09 -0700
Subject: [PATCH 19/36] fix(executorch): name both TensorRT versions when an
 engine cannot load

A TensorRT engine loads only on the TensorRT version that built it,
unless it was built version compatible. When a program exported with one
version is loaded with another, TensorRT refuses the engine, but its
error names internal serialization numbers, not versions. An engine
built with TensorRT 11.4 and loaded with 11.3 gives:

  Version tag does not match. Note: Current Version: 244, Serialized
  Engine Version: 245

So the user cannot tell which TensorRT to install.

Export now records its TensorRT version in the delegate metadata, as
tensorrt_version. When TensorRT cannot read an engine, the delegate logs
both versions:

  TensorRTBackend::init: the engine was exported with TensorRT
  11.4.0.106, and this runtime is TensorRT 11.3.0

For a program exported before this change, it names the runtime version
and says the program does not record the other one. The version is one
more key in the metadata JSON, so the header and its magic stay the
same, and a delegate that does not know the key ignores it.

The new runtime test zeroes the start of a stored engine, so TensorRT
refuses it on the same path, and checks that the log names both
versions. It does not install a second TensorRT version. The
serialization and preprocess tests check that the version is written and
read back.
---
 .../executorch/TensorRTBlobHeader.h           |  1 +
 cpp/src/torch_tensorrt/executorch/README.md   |  4 ++
 .../executorch/TensorRTBackend.cpp            | 41 +++++++++++++++++--
 .../executorch/TensorRTBlobHeader.cpp         |  7 ++++
 py/torch_tensorrt/executorch/backend.py       |  2 +
 py/torch_tensorrt/executorch/serialization.py |  5 +++
 tests/py/dynamo/executorch/test_backend.py    |  2 +
 .../py/dynamo/executorch/test_load_errors.py  | 30 ++++++++++++++
 .../dynamo/executorch/test_serialization.py   |  2 +
 9 files changed, 90 insertions(+), 4 deletions(-)

diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h
index 8c9c31c33d5..216f48b2991 100644
--- a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h
+++ b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h
@@ -51,6 +51,7 @@ struct TensorRTBlobHeader {
   bool hardware_compatible = false; ///< Engine built hardware-compatible; false when unstated.
   int device_id = 0; ///< Device the engine was built for; 0 when unstated.
   std::string engine_key; ///< Named-data key of an engine stored outside the blob; empty when it is inline.
+  std::string tensorrt_version; ///< TensorRT version that exported the engine; empty when the blob does not say.
 
   /**
    * @brief Start of the serialized engine inside a blob.
diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md
index 125f28a73d8..fda72dbe62f 100644
--- a/cpp/src/torch_tensorrt/executorch/README.md
+++ b/cpp/src/torch_tensorrt/executorch/README.md
@@ -201,6 +201,10 @@ but corrupted engine bytes can load and return wrong answers.
   deserializes the engine, applies the weight streaming budget, or creates the execution context.
   The log then shows the free and total GPU memory. Other TensorRT load failures return
   `Error::InvalidProgram`, with TensorRT's reason in the log.
+- An engine loads only on the TensorRT version that built it, unless it was built version
+  compatible. When TensorRT cannot read an engine, the log names the TensorRT version that exported
+  it and the one the delegate runs on. A program exported before the export side recorded that
+  version says so instead.
 - An engine saved without its weights, for example compiled with
   `immutable_weights=False, strip_engine_weights=True`, is refused with `Error::InvalidProgram`,
   because this delegate cannot refit it. To check, the delegate asks TensorRT how many weight bytes
diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
index c5e49f3679b..4be9d71a48c 100644
--- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
+++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp
@@ -130,6 +130,32 @@ Error report_load_error(const char* operation) {
   return out_of_memory ? Error::MemoryAllocationFailed : Error::InvalidProgram;
 }
 
+// TensorRT reads an engine only with the version that built it, unless it was built version
+// compatible, and its own error then names internal serialization tags rather than versions.
+void log_tensorrt_versions(const std::string& exported_version) {
+  // TensorRT 10 and later encode the version as major * 10000 + minor * 100 + patch.
+  const int32_t version = getInferLibVersion();
+  const int major = version / 10000;
+  const int minor = version / 100 % 100;
+  const int patch = version % 100;
+  if (exported_version.empty()) {
+    ET_LOG(
+        Error,
+        "TensorRTBackend::init: this runtime is TensorRT %d.%d.%d, and the program does not record which TensorRT exported the engine",
+        major,
+        minor,
+        patch);
+  } else {
+    ET_LOG(
+        Error,
+        "TensorRTBackend::init: the engine was exported with TensorRT %s, and this runtime is TensorRT %d.%d.%d",
+        exported_version.c_str(),
+        major,
+        minor,
+        patch);
+  }
+}
+
 Error check_registration() {
   if (kRegistrationResult != Error::Ok) {
     ET_LOG(Error, "TensorRTBackend registration failed: %s", ::executorch::runtime::to_string(kRegistrationResult));
@@ -410,11 +436,16 @@ Error load_engine(
     const void* data,
     uint64_t size,
     const WsBudget& request,
+    const std::string& exported_version,
     std::shared_ptr<nvinfer1::ICudaEngine>& out) {
   last_load_error = {};
   std::shared_ptr<nvinfer1::ICudaEngine> engine(runtime.deserializeCudaEngine(data, size), TRTDeleter{});
   if (engine == nullptr) {
-    return report_load_error("deserialize TensorRT engine");
+    const Error err = report_load_error("deserialize TensorRT engine");
+    if (err != Error::MemoryAllocationFailed) {
+      log_tensorrt_versions(exported_version);
+    }
+    return err;
   }
   // A plan serialized without its weights (built with STRIP_PLAN, or saved with EXCLUDE_WEIGHTS)
   // deserializes and runs without an error, but computes with placeholder weights, and this backend
@@ -464,6 +495,7 @@ Error acquire_shared_engine(
     uint64_t size,
     int device_id,
     const WsBudget& request,
+    const std::string& exported_version,
     std::shared_ptr<nvinfer1::ICudaEngine>& out) {
   const SharedEngineKey key = shared_engine_key(data, size, device_id, request.valid ? request.bytes : -1);
   SharedEngines& engines = shared_engines();
@@ -476,7 +508,7 @@ Error acquire_shared_engine(
   // Declared out here so a copy that lost the race is freed after the lock is released.
   std::shared_ptr<nvinfer1::ICudaEngine> loaded;
   if (out == nullptr) {
-    const Error err = load_engine(runtime, data, size, request, loaded);
+    const Error err = load_engine(runtime, data, size, request, exported_version, loaded);
     if (err != Error::Ok) {
       // A racing load of the same bytes may have published its engine meanwhile, and this load may
       // have failed only because the device had no room for a second copy.
@@ -1422,9 +1454,10 @@ Result<DelegateHandle*> TensorRTBackend::init(
   {
     const std::shared_lock<std::shared_mutex> capture_lock(cuda_graph_capture_mutex());
     if (share) {
-      err = acquire_shared_engine(*runtime, engine_data, engine_size, handle->device_id, ws_request, handle->engine);
+      err = acquire_shared_engine(
+          *runtime, engine_data, engine_size, handle->device_id, ws_request, header.tensorrt_version, handle->engine);
     } else {
-      err = load_engine(*runtime, engine_data, engine_size, ws_request, handle->engine);
+      err = load_engine(*runtime, engine_data, engine_size, ws_request, header.tensorrt_version, handle->engine);
     }
     // The engine now lives on the GPU, so release the host copy before building the context.
     named_engine.reset();
diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp
index a5e2aff3a27..6a9484e52c6 100644
--- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp
+++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp
@@ -487,6 +487,13 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso
       return false;
     }
   }
+  // Only an error message reads it, so a value that is not a string is ignored rather than refused.
+  out.tensorrt_version.clear();
+  const std::size_t version_key = find_top_level_key(json, "\"tensorrt_version\"");
+  if (version_key != std::string::npos &&
+      parse_string(json, skip_ws(json, json.find(':', version_key) + 1), out.tensorrt_version) == std::string::npos) {
+    out.tensorrt_version.clear();
+  }
   const std::size_t hw_key = find_top_level_key(json, "\"hardware_compatible\"");
   const std::size_t device_key = find_top_level_key(json, "\"device_id\"");
   return parse_bool_after_key(json, hw_key, "\"hardware_compatible\"", out.hardware_compatible) &&
diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py
index 6f08b8bf047..ac8bfc67a19 100644
--- a/py/torch_tensorrt/executorch/backend.py
+++ b/py/torch_tensorrt/executorch/backend.py
@@ -10,6 +10,7 @@
 import operator
 from typing import Any, Container, Iterable, List, Optional, Set, final
 
+import tensorrt
 import torch
 import torch.fx
 from executorch.exir._serialize._named_data_store import NamedDataStore
@@ -662,6 +663,7 @@ def preprocess(
             device_id=_parse_device_id(engine_info[DEVICE_IDX]),
             serialized_metadata=_get_str(engine_info, SERIALIZED_METADATA_IDX),
             target_platform=_get_str(engine_info, TARGET_PLATFORM_IDX),
+            tensorrt_version=tensorrt.__version__,
             # Content-keyed, so identical engines in different methods share one
             # entry and different engines can never collide on a key.
             engine_key="tensorrt_engine_" + hashlib.sha256(engine_bytes).hexdigest(),
diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py
index dc7b0a56814..d0214fd0cd3 100644
--- a/py/torch_tensorrt/executorch/serialization.py
+++ b/py/torch_tensorrt/executorch/serialization.py
@@ -64,6 +64,9 @@ class TensorRTBlobMetadata:
     # Named-data key of the engine when it is stored outside the blob, which then
     # carries an empty engine. Empty means the engine follows the metadata inline.
     engine_key: str = ""
+    # TensorRT version that exported the engine, so a runtime whose TensorRT cannot
+    # read the engine can name both versions. Empty when unknown.
+    tensorrt_version: str = ""
 
     def to_json(self) -> bytes:
         # Keep field order stable because the C++ parser is intentionally small.
@@ -92,6 +95,7 @@ def to_json(self) -> bytes:
             "device_id": self.device_id,
             "serialized_metadata": self.serialized_metadata,
             "target_platform": self.target_platform,
+            "tensorrt_version": self.tensorrt_version,
         }
         if self.engine_key:
             data["engine_key"] = self.engine_key
@@ -119,6 +123,7 @@ def from_json(cls, data: bytes) -> "TensorRTBlobMetadata":
             serialized_metadata=parsed.get("serialized_metadata", ""),
             target_platform=parsed.get("target_platform", ""),
             engine_key=parsed.get("engine_key", ""),
+            tensorrt_version=parsed.get("tensorrt_version", ""),
         )
 
 
diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py
index 4cafbdb1ffe..580cbfdb411 100644
--- a/tests/py/dynamo/executorch/test_backend.py
+++ b/tests/py/dynamo/executorch/test_backend.py
@@ -11,6 +11,7 @@
 
 executorch = pytest.importorskip("executorch.exir")
 
+import tensorrt  # noqa: E402
 import torch  # noqa: E402
 from executorch.exir.backend.compile_spec_schema import CompileSpec  # noqa: E402
 from torch.export.graph_signature import InputKind  # noqa: E402
@@ -201,6 +202,7 @@ def test_preprocess_serializes_engine_blob():
     assert _named_engine(result) == b"engine-bytes"
     _, metadata = deserialize_engine(result.processed_bytes)
     assert metadata.device_id == 2
+    assert metadata.tensorrt_version == tensorrt.__version__
     assert [binding.name for binding in metadata.io_bindings] == ["x", "y"]
     assert [binding.is_input for binding in metadata.io_bindings] == [True, False]
 
diff --git a/tests/py/dynamo/executorch/test_load_errors.py b/tests/py/dynamo/executorch/test_load_errors.py
index 8ae74939470..71cfc9f9cb5 100644
--- a/tests/py/dynamo/executorch/test_load_errors.py
+++ b/tests/py/dynamo/executorch/test_load_errors.py
@@ -77,3 +77,33 @@ def test_loading_with_no_visible_gpu_fails_with_a_clear_error(tmp_path):
     )
     assert result.returncode == 3, (result.returncode, result.stderr[-2000:])
     assert "CUDA is unavailable" in result.stdout + result.stderr
+
+
+def test_an_engine_tensorrt_cannot_read_names_both_tensorrt_versions(tmp_path, capfd):
+    """A real version mismatch needs a second TensorRT install, so this corrupts the stored engine.
+
+    TensorRT then refuses it on the same path it takes for an engine from another version, which
+    is the path that has to name both versions.
+    """
+    import tensorrt
+    import torch_tensorrt_executorch_runtime  # noqa: F401
+    from executorch.exir._serialize._program import deserialize_pte_binary
+    from executorch.runtime import Runtime
+
+    pte = tmp_path / "model.pte"
+    _save(pte)
+    data = bytearray(pte.read_bytes())
+    named = deserialize_pte_binary(bytes(data)).named_data
+    (engine,) = [bytes(buffer) for buffer in named.buffers]
+    start = data.find(engine)
+    assert start > 0 and data.find(engine, start + 1) == -1
+    data[start : start + 64] = bytes(64)
+    pte.write_bytes(data)
+
+    program = Runtime.get().load_program(pte)
+    with pytest.raises(RuntimeError, match="0x:?23") as raised:
+        program.load_method("forward")
+    log = str(raised.value) + capfd.readouterr().err
+    runtime_version = ".".join(tensorrt.__version__.split(".")[:3])
+    assert f"exported with TensorRT {tensorrt.__version__}" in log, log
+    assert f"this runtime is TensorRT {runtime_version}" in log, log
diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py
index f2f8ef9514f..c933d6f6161 100644
--- a/tests/py/dynamo/executorch/test_serialization.py
+++ b/tests/py/dynamo/executorch/test_serialization.py
@@ -25,6 +25,7 @@ def test_serialize_engine_writes_tr01_blob():
         ],
         device_id=1,
         hardware_compatible=True,
+        tensorrt_version="11.4.0.106",
     )
 
     blob = serialize_engine(b"engine-bytes", metadata)
@@ -36,6 +37,7 @@ def test_serialize_engine_writes_tr01_blob():
     assert engine == b"engine-bytes"
     assert parsed.device_id == 1
     assert parsed.hardware_compatible is True
+    assert parsed.tensorrt_version == "11.4.0.106"
     assert [b.name for b in parsed.io_bindings] == ["x", "y"]
     assert [b.is_input for b in parsed.io_bindings] == [True, False]
 

From e9f3b49c93eb2e0ed11511e01d86cd5f1bd113a7 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:46:26 -0700
Subject: [PATCH 20/36] fix(conversion): keep a static dim when expand
 broadcasts to a runtime size

`impl.slice.expand` treated every static dim whose target size is only known
at runtime as a broadcast. That is only right for a dim of size 1. For a
larger dim it repeated the first element instead. For example, expanding a
`(4, 8)` tensor to `(n, 8)`, where `n` is the True count of a mask, returned
four copies of row 0 even when `n` was 4.

PyTorch only lets a dim larger than 1 expand to its own size, so expand now
keeps such a dim. Every caller gets the fix. The index_put converter no
longer needs its own `_broadcast_target` helper for this, so it passes the
target shape straight to expand again, and keeps only the shape check it had
before. Static shapes behave as they did before the data-dependent mask fix.

Test: expand of a static `(4, 8)` tensor to the True count of a mask, compared
with eager. It fails without the fix. The data-dependent index_put test now
also checks that a True count that differs from the value rows fails at
runtime, as in PyTorch.
---
 .../dynamo/conversion/impl/select.py          | 42 +++++++------------
 .../dynamo/conversion/impl/slice/ops.py       |  8 ++++
 .../py/dynamo/conversion/test_expand_aten.py  | 18 ++++++++
 .../dynamo/conversion/test_index_put_aten.py  |  5 +++
 4 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py
index 37e2e45c71d..aaab9fea6db 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/select.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py
@@ -683,29 +683,6 @@ def index_put_scatter_add_plugin(
     return out
 
 
-def _broadcast_target(
-    expected_shape: Sequence[Union[int, ITensor]], tensor_shape: Sequence[int]
-) -> List[Union[int, ITensor]]:
-    """The shape to expand a tensor laid out as (N, ...) to, broadcasting like PyTorch.
-
-    expand broadcasts any static dim whose target is an ITensor, which is only right
-    for size 1. A data-dependent N, such as the True count of a boolean mask, must
-    match a larger static first dim at runtime in PyTorch, so keep that dim.
-    """
-    tensor_shape = [1] * (len(expected_shape) - len(tensor_shape)) + list(tensor_shape)
-    shape: List[Union[int, ITensor]] = []
-    for k, (exp_dim, dim) in enumerate(zip(expected_shape, tensor_shape)):
-        if k == 0 and isinstance(exp_dim, ITensor) and dim not in (1, DYNAMIC_DIM):
-            shape.append(dim)
-        elif dim in (1, DYNAMIC_DIM) or exp_dim == dim:
-            shape.append(exp_dim)
-        else:
-            raise ValueError(
-                f"Cannot broadcast {tuple(tensor_shape)} to {expected_shape}"
-            )
-    return shape
-
-
 def index_put_converter(
     ctx: ConversionContext,
     target: Target,
@@ -876,7 +853,7 @@ def index_put_converter(
             source_ir,
             f"{name}_expand_idx_I_{i}",
             idx_reshaped,
-            tuple(_broadcast_target((N, F_volume), idx_reshaped.shape)),
+            (N, F_volume),
         )
         I_tensors.append(expanded_idx)
 
@@ -1022,7 +999,6 @@ def index_put_converter(
 
     # Define the expected shape based on (N,) + F_shape_values.
     # N may be a ITensor when it comes from a dynamic source (e.g. nonzero).
-    # _broadcast_target decides which values dims expand to it.
     expected_shape = (N,) + tuple(F_shape_values)
 
     # Broadcast 'values' to match the expected shape
@@ -1059,7 +1035,7 @@ def index_put_converter(
                 source_ir,
                 f"{name}_expand_values",
                 values_permuted,
-                tuple(_broadcast_target(expected_shape, values_permuted.shape)),
+                expected_shape,
             )
         elif K == 1 and len(values.shape) == 1 and F and I[0] > max(F):
             # For data[..., idx] = values, a 1-D values tensor carries the
@@ -1136,13 +1112,25 @@ def index_put_converter(
             )
         else:
             # Discontinuous case (K > 1 or K == 0)
+            values_shape_padded = [1] * (
+                len(expected_shape) - len(values.shape)
+            ) + list(values.shape)
+            for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
+                if not (
+                    isinstance(exp_dim, ITensor)
+                    or DYNAMIC_DIM in (exp_dim, val_dim)
+                    or val_dim in (1, exp_dim)
+                ):
+                    raise ValueError(
+                        f"Cannot broadcast {values.shape} to {expected_shape}"
+                    )
             values_expanded = impl.slice.expand(
                 ctx,
                 target,
                 source_ir,
                 f"{name}_expand_values",
                 values,
-                tuple(_broadcast_target(expected_shape, values.shape)),
+                expected_shape,
             )
 
     # Flatten values to (N * F_volume,)
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
index 86a90ebd1ed..32a01230d6b 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
@@ -310,6 +310,14 @@ def expand(
     # After the above padding, the shape and tensor rank must be equal
     assert len(input_t.shape) == shape_rank
 
+    # PyTorch only expands a dim larger than 1 to its own size, so a runtime
+    # target there must equal the static input dim. Keep the dim instead of
+    # broadcasting its first element.
+    shape = [
+        i if isinstance(o, TRTTensor) and i not in (1, DYNAMIC_DIM) else o
+        for i, o in zip(input_t.shape, shape)
+    ]
+
     # Configure the start, strides and output shape tensors
     start = tuple([0] * shape_rank)
 
diff --git a/tests/py/dynamo/conversion/test_expand_aten.py b/tests/py/dynamo/conversion/test_expand_aten.py
index fed09e29671..fe0b5a6e7b0 100644
--- a/tests/py/dynamo/conversion/test_expand_aten.py
+++ b/tests/py/dynamo/conversion/test_expand_aten.py
@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
@@ -84,6 +85,23 @@ def forward(self, x):
             ExpandTargetDynamic(), input_specs, use_dynamo_tracer=True
         )
 
+    def test_expand_static_dim_to_data_dependent_size(self):
+        # The True count is a runtime target that must equal the static dim of 4.
+        class ExpandToCount(nn.Module):
+            def forward(self, x, mask):
+                return x.expand(mask.nonzero().shape[0], 8)
+
+        mask = torch.zeros(10, dtype=torch.bool, device="cuda")
+        mask[3:7] = True
+        inputs = (torch.randn(4, 8, device="cuda"), mask)
+        compiled = torch_tensorrt.dynamo.compile(
+            torch.export.export(ExpandToCount(), inputs),
+            arg_inputs=inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        torch.testing.assert_close(compiled(*inputs), ExpandToCount()(*inputs))
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py
index 4ccf1747fbf..14d24de3449 100644
--- a/tests/py/dynamo/conversion/test_index_put_aten.py
+++ b/tests/py/dynamo/conversion/test_index_put_aten.py
@@ -556,6 +556,11 @@ def forward(self, data, mask, values):
                     pass_through_build_failures=True,
                 )
                 torch.testing.assert_close(compiled(*inputs), model(*inputs))
+                # A True count that differs from the value rows fails, as in PyTorch.
+                fewer = case_mask.clone()
+                fewer[..., 6] = False
+                with self.assertRaises(RuntimeError):
+                    compiled(data, fewer, values)
 
     def test_updates_are_cast_to_int64_destination_dtype(self):
         class IndexPutArgmax(torch.nn.Module):

From a1c037dfe1a826129b3ea67f48e6935e74c1a2c5 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:09:46 -0700
Subject: [PATCH 21/36] fix(conversion): broadcast a one-entry boolean mask
 across static indices in index_put

`y[mask, idx] = values`, with a static `idx` of 4 column indices, works in
PyTorch when `mask` has 4 True entries and also when it has exactly 1: index
tensors broadcast together, so the one selected row is paired with each of
the 4 columns. The compiled engine matched eager for 4 True entries, but for
1 it failed at runtime with "reshape would change volume 32 to 8".

The index_put converter took the number of updated rows N from the mask's
True count whenever an index had a runtime length. But a runtime length must
be 1 or equal a static length larger than 1, so N is that static length. The
converter now uses it.

Expanding the mask's indices to that static N then needs `impl.slice.expand`
to broadcast a runtime dim of size 1 to a static size. Before, expand assumed
a runtime dim already had the target size. It now picks the stride at
runtime: 0 when the dim is 1, and 1 otherwise. Only this case pays for the
two extra shape layers.

Test: the data-dependent index_put test now also runs the column case with a
single True entry and compares with eager. It fails without the fix.
---
 .../dynamo/conversion/impl/select.py          | 13 ++++++------
 .../dynamo/conversion/impl/slice/ops.py       | 13 ++++++++++--
 .../dynamo/conversion/test_index_put_aten.py  | 21 ++++++++++++++-----
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py
index aaab9fea6db..e0b5c4f6e82 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/select.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py
@@ -789,13 +789,12 @@ def index_put_converter(
                             0,
                         )
                     )
-        # When any index has a dynamic size, use the first dynamic value
-        # (all valid indices are guaranteed to have the same N after broadcasting).
-        # Python's max() cannot compare ITensors, so we avoid it when dynamic.
-        if any(isinstance(s, ITensor) for s in index_shapes):
-            N = next(s for s in index_shapes if isinstance(s, ITensor))
-        else:
-            N = max(index_shapes) if index_shapes else 1
+        # Index lengths broadcast together. A dynamic length must be 1 or equal a
+        # static length larger than 1, so that static length is N. Otherwise use
+        # the first dynamic length, since Python's max() cannot compare ITensors.
+        static_n = max((s for s in index_shapes if isinstance(s, int)), default=1)
+        dynamic_n = [s for s in index_shapes if isinstance(s, ITensor)]
+        N = dynamic_n[0] if dynamic_n and static_n == 1 else static_n
     else:
         N = 1
 
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
index 32a01230d6b..eab48b9cd00 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
@@ -325,7 +325,7 @@ def expand(
     # stride should be 1 for all non-broadcasted dims
     stride = []
     input_tensor_shape = tuple(input_t.shape)
-    for i, o in zip(input_tensor_shape, shape):
+    for k, (i, o) in enumerate(zip(input_tensor_shape, shape)):
         # If input dim and target shape dim are static, broadcast if they are not equal
         # If input dim is known and target shape dim is dynamic we treat it as a broadcasted dim
         if (
@@ -337,6 +337,15 @@ def expand(
             stride.append(int(i == o))
         elif isinstance(i, int) and i != DYNAMIC_DIM and isinstance(o, TRTTensor):
             stride.append(0)
+        elif i == DYNAMIC_DIM and isinstance(o, int) and o != DYNAMIC_DIM:
+            # A runtime dim is either 1, and broadcasts, or already the target size.
+            size = get_shape(
+                ctx, target, source_ir, f"{name}_stride_dim{k}", input_t, k
+            )
+            is_full = impl.elementwise.gt(
+                ctx, target, source_ir, f"{name}_stride_gt{k}", size, 1
+            )
+            stride.append(cast_trt_tensor(ctx, is_full, trt.int32, f"{name}_stride{k}"))
         else:
             # No broadcasting is happening. The output should have the same size as input at this dimension.
             stride.append(1)
@@ -354,7 +363,7 @@ def expand(
 
     target_shape_t = target_shape
     # Handle dynamic shapes case where shape has dynamic dimension
-    if any(isinstance(ele, TRTTensor) for ele in target_shape_t):
+    if any(isinstance(ele, TRTTensor) for ele in (*target_shape_t, *stride)):
         target_shape_t = cat(
             ctx,
             target,
diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py
index 14d24de3449..02f38a7c2f9 100644
--- a/tests/py/dynamo/conversion/test_index_put_aten.py
+++ b/tests/py/dynamo/conversion/test_index_put_aten.py
@@ -541,12 +541,20 @@ def forward(self, data, mask, values):
         values = torch.randn(4, 8, device="cuda")
         mask = torch.zeros(1, 10, dtype=torch.bool, device="cuda")
         mask[0, 3:7] = True
+        one_true = torch.zeros(10, dtype=torch.bool, device="cuda")
+        one_true[5] = True
         cases = [
-            (AddRows(), torch.randn(1, 10, 8, device="cuda"), mask),
-            (SetRows(), torch.randn(10, 8, device="cuda"), mask[0]),
-            (SetRowsAtColumns(), torch.randn(10, 6, 8, device="cuda"), mask[0]),
+            (AddRows(), torch.randn(1, 10, 8, device="cuda"), mask, ()),
+            (SetRows(), torch.randn(10, 8, device="cuda"), mask[0], ()),
+            # One True entry broadcasts across the four column indices.
+            (
+                SetRowsAtColumns(),
+                torch.randn(10, 6, 8, device="cuda"),
+                mask[0],
+                (one_true,),
+            ),
         ]
-        for model, data, case_mask in cases:
+        for model, data, case_mask, other_masks in cases:
             with self.subTest(type(model).__name__):
                 inputs = (data, case_mask, values)
                 compiled = torchtrt.dynamo.compile(
@@ -555,7 +563,10 @@ def forward(self, data, mask, values):
                     min_block_size=1,
                     pass_through_build_failures=True,
                 )
-                torch.testing.assert_close(compiled(*inputs), model(*inputs))
+                for run_mask in (case_mask, *other_masks):
+                    torch.testing.assert_close(
+                        compiled(data, run_mask, values), model(data, run_mask, values)
+                    )
                 # A True count that differs from the value rows fails, as in PyTorch.
                 fewer = case_mask.clone()
                 fewer[..., 6] = False

From ad67feba539b3da9345dd7f8a1e40dc90d40c8c6 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:06:11 -0700
Subject: [PATCH 22/36] fix(conversion): broadcast the source of aten.copy to
 the destination shape

`aten.copy(dest, src)` returns a tensor with the shape of `dest`, and `src`
only has to broadcast to it. The converter cast `src` and returned it as it
was, so copying a `(3,)` tensor into a `(2, 3)` buffer returned shape `(3,)`
with no error. The same happened with dynamic shapes, for example a `(1, 5)`
source copied into a `(batch, 5)` destination.

With `truncate_double=True` this now also reached float64 buffers, which used
to fail the build, so a copy into a float64 buffer returned a wrong shape
silently instead of failing.

When the shapes differ, the converter now expands `src` to the shape of
`dest` before the cast, reading each dynamic size of `dest` at run time.

Test: static and dynamic broadcast cases in the copy tests, and the float64
buffer copy in the cast tests now broadcasts its source. All of them fail
without the fix and match eager with it.
---
 .../dynamo/conversion/aten_ops_converters.py  | 19 +++++++++++++++++++
 tests/py/dynamo/conversion/test_casts.py      |  5 +++--
 tests/py/dynamo/conversion/test_copy_aten.py  | 19 +++++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 49d0fde4503..e1430eb8122 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -4826,6 +4826,25 @@ def aten_ops_copy(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     dest, src = args[0], args[1]
+    # copy broadcasts src, so the result always has the shape of dest.
+    if tuple(src.shape) != tuple(dest.shape):
+        src = impl.slice.expand(
+            ctx,
+            target,
+            SourceIR.ATEN,
+            f"{name}_expand",
+            src,
+            [
+                (
+                    impl.shape.shape(
+                        ctx, target, SourceIR.ATEN, f"{name}_dest_dim{i}", dest, i
+                    )
+                    if size == DYNAMIC_DIM
+                    else size
+                )
+                for i, size in enumerate(dest.shape)
+            ],
+        )
     return impl.cast.to_copy(
         ctx,
         target,
diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py
index ac2926ee2cb..cd8777c1010 100644
--- a/tests/py/dynamo/conversion/test_casts.py
+++ b/tests/py/dynamo/conversion/test_casts.py
@@ -102,11 +102,12 @@ def forward(self, x):
                 y = x * 2
                 return y, torch.ops.aten._to_copy.default(y, dtype=torch.float64)
 
-        # aten.copy reaches the same cast without the _to_copy validator.
+        # aten.copy reaches the same cast without the _to_copy validator, and
+        # broadcasts its source to the buffer shape.
         class CopyIntoDouble(nn.Module):
             def __init__(self):
                 super().__init__()
-                self.register_buffer("dest", torch.zeros(1, 3, 10, dtype=torch.double))
+                self.register_buffer("dest", torch.zeros(2, 3, 10, dtype=torch.double))
 
             def forward(self, x):
                 return torch.ops.aten.copy.default(self.dest, x * 2)
diff --git a/tests/py/dynamo/conversion/test_copy_aten.py b/tests/py/dynamo/conversion/test_copy_aten.py
index 2752e1faf14..5ac56d31d12 100644
--- a/tests/py/dynamo/conversion/test_copy_aten.py
+++ b/tests/py/dynamo/conversion/test_copy_aten.py
@@ -17,6 +17,8 @@ class TestCopyConverter(DispatchTestCase):
             ((1, 10), (1, 10), False),
             ((2, 3, 4), (2, 3, 4), True),
             ((2, 3, 4, 5), (2, 3, 4, 5), True),
+            ((2, 3), (3,), False),
+            ((2, 3, 4), (2, 1, 4), False),
         ]
     )
     def test_copy_float(self, input_shape, src_shape, non_blocking):
@@ -91,6 +93,23 @@ def forward(self, input, src):
             input_specs,
         )
 
+    @parameterized.expand([((5,),), ((1, 5),)])
+    def test_dynamic_shape_copy_broadcast(self, src_shape):
+        class Copy(nn.Module):
+            def forward(self, input, src):
+                return torch.ops.aten.copy.default(input, src)
+
+        input_specs = [
+            Input(
+                min_shape=(1, 5),
+                opt_shape=(3, 5),
+                max_shape=(8, 5),
+                dtype=torch.float32,
+            ),
+            Input(shape=src_shape, dtype=torch.float32),
+        ]
+        self.run_test_with_dynamic_shape(Copy(), input_specs)
+
     def test_copy_mixed_dtype_dest_fp32_src_fp16(self):
         # Regression for pytorch/TensorRT#4265: aten.copy(self, src) returns a
         # tensor with self's dtype (src is implicitly cast). The previous

From 8648760fb97b02ca1fbba023a159665c74369014 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:12:35 -0700
Subject: [PATCH 23/36] fix(conversion): check float64 aten.copy with the same
 rule as _to_copy

With `truncate_double=True`, a float64 `_to_copy` is checked before it goes
to TensorRT: the source has to be a type float32 holds exactly, and the result
has to be contiguous, because the TensorRT copy always is. `aten.copy` into a
float64 buffer had no such check, so it reached the cast with no layout check
at all. A copy into a channels_last float64 buffer built and returned a
contiguous tensor where PyTorch keeps the layout of the buffer, and a copy of
an int64 source failed the whole build.

The two checks also used different lists of source types: the `_to_copy`
check accepted only float sources, while the cast itself also accepted bool
and 8-bit integers. So `x.to(torch.float64)` of a bool tensor stayed in
PyTorch and split the graph, while the same values copied into a float64
buffer converted.

Now one list of source types lives next to the cast, and both the cast and the
checks read it. `aten.copy` gets a check that applies the `_to_copy` rule when
the result is float64, so those copies convert or stay in PyTorch the same
way. Copies to other types are not affected.

Test: a channels_last float64 copy keeps its layout, an int64 copy into a
float64 buffer stays in PyTorch, and a bool to float64 cast converts, each in
one engine and matching eager. All three fail without the fix.
---
 .../dynamo/conversion/aten_ops_converters.py  | 40 ++++++++++++-------
 .../dynamo/conversion/impl/cast.py            | 21 ++++++++--
 tests/py/dynamo/conversion/test_casts.py      | 40 +++++++++++++++++++
 3 files changed, 83 insertions(+), 18 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index e1430eb8122..dec408c2fb4 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2162,21 +2162,18 @@ def aten_ops_permute(
     )
 
 
-def _truncatable_float64_copy(node: Node) -> bool:
+def _truncatable_float64_copy(node: Node, source: Argument) -> bool:
     """Whether a float64 copy can run as a float32 copy and give PyTorch's result.
 
-    The source has to be a float type TensorRT casts to float32, and the PyTorch
-    result has to have the contiguous layout a TensorRT copy produces. The check
-    does not depend on where the copy sits, because partitioning can turn any copy
-    into an engine output and the converter is validated again then.
+    The source has to be a type float32 holds exactly, the same rule impl.cast.to_copy
+    applies, and the PyTorch result has to have the contiguous layout a TensorRT copy
+    produces. The check does not depend on where the copy sits, because partitioning
+    can turn any copy into an engine output and the converter is validated again then.
     """
-    input_node = node.args[0] if node.args else None
-    input_meta = input_node.meta.get("val") if isinstance(input_node, Node) else None
-    if not isinstance(input_meta, torch.Tensor) or input_meta.dtype not in (
-        torch.float16,
-        torch.bfloat16,
-        torch.float32,
-        torch.float64,
+    input_meta = source.meta.get("val") if isinstance(source, Node) else None
+    if (
+        not isinstance(input_meta, torch.Tensor)
+        or input_meta.dtype not in impl.cast.FLOAT64_TRUNCATABLE_SOURCES
     ):
         return False
     # Alone on a placeholder after partitioning, a dynamic copy must also pass
@@ -2246,7 +2243,7 @@ def validate_dtype(
                 to_copy_node.kwargs["dtype"] is torch.float64
                 and settings is not None
                 and settings.truncate_double
-                and _truncatable_float64_copy(to_copy_node)
+                and _truncatable_float64_copy(to_copy_node, to_copy_node.args[0])
             ):
                 return True
             else:
@@ -4812,7 +4809,22 @@ def aten_ops_trunc(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.copy.default, supports_dynamic_shapes=True)
+def _copy_validator(node: Node, settings: Optional[CompilationSettings] = None) -> bool:
+    output_meta = node.meta.get("val")
+    if not isinstance(output_meta, torch.Tensor) or output_meta.dtype != torch.float64:
+        return True
+    return (
+        settings is not None
+        and settings.truncate_double
+        and _truncatable_float64_copy(node, node.args[1])
+    )
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.copy.default,
+    capability_validator=_copy_validator,
+    supports_dynamic_shapes=True,
+)
 @enforce_tensor_types(
     {
         1: (TRTTensor,),
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/cast.py b/py/torch_tensorrt/dynamo/conversion/impl/cast.py
index 1f373348007..478f274c19f 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/cast.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/cast.py
@@ -21,6 +21,20 @@
 
 LOGGER: logging.Logger = logging.getLogger(__name__)
 
+# With truncate_double, a float64 result from these sources runs in float32, which
+# holds their values exactly. A float64 source is already float32 inside TensorRT.
+FLOAT64_TRUNCATABLE_SOURCES = frozenset(
+    {
+        torch.bool,
+        torch.int8,
+        torch.uint8,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    }
+)
+
 
 def to_copy(
     ctx: ConversionContext,
@@ -34,13 +48,12 @@ def to_copy(
     if not isinstance(input, TRTTensor):
         input = get_trt_tensor(ctx, input, f"{name}_copy_tensor")
 
-    # TensorRT has no float64. With truncate_double the caller accepts float32,
-    # which holds these source types exactly but not every larger integer.
+    # TensorRT has no float64. With truncate_double the caller accepts float32.
     if (
         ctx.compilation_settings.truncate_double
         and _enums.dtype._from(dtype) == _enums.dtype.f64
-        and input.dtype
-        in {trt.bool, trt.int8, trt.uint8, trt.float16, trt.bfloat16, trt.float32}
+        and _enums.dtype._from(input.dtype).to(torch.dtype)
+        in FLOAT64_TRUNCATABLE_SOURCES
     ):
         dtype = torch.float32
 
diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py
index cd8777c1010..d27d2aaa753 100644
--- a/tests/py/dynamo/conversion/test_casts.py
+++ b/tests/py/dynamo/conversion/test_casts.py
@@ -165,6 +165,46 @@ def forward(self, x):
                     x = torch.rand(n, 5, device="cuda") + 0.5
                     torch.testing.assert_close(compiled(x), model(x))
 
+    def test_float64_copies_share_one_truncation_rule(self):
+        # The float64 rule for _to_copy applies to aten.copy too: a source float32
+        # cannot hold exactly, or a result layout TensorRT cannot produce, stays in
+        # PyTorch, and the other cases convert.
+        class CopyIntoDouble(nn.Module):
+            def __init__(self, memory_format):
+                super().__init__()
+                dest = torch.zeros(2, 3, 4, 5, dtype=torch.double)
+                self.register_buffer("dest", dest.to(memory_format=memory_format))
+
+            def forward(self, x):
+                return torch.ops.aten.copy.default(self.dest, x * 2)
+
+        class BoolToDouble(nn.Module):
+            def forward(self, x):
+                return (x > 0.5).to(torch.float64) * 3
+
+        cases = (
+            (CopyIntoDouble(torch.channels_last), torch.rand(2, 3, 4, 5)),
+            (CopyIntoDouble(torch.contiguous_format), torch.randint(9, (2, 3, 4, 5))),
+            (BoolToDouble(), torch.rand(2, 3, 4, 5)),
+        )
+        for model, x in cases:
+            with self.subTest(type(model).__name__, dtype=x.dtype):
+                model = model.eval().cuda()
+                x = x.cuda()
+                compiled = torch_tensorrt.dynamo.compile(
+                    torch.export.export(model, (x,)),
+                    arg_inputs=[x],
+                    min_block_size=1,
+                    truncate_double=True,
+                )
+                self.assertEqual(
+                    sum("_run_on_acc" in name for name, _ in compiled.named_children()),
+                    1,
+                )
+                out, expected = compiled(x), model(x)
+                torch.testing.assert_close(out, expected)
+                self.assertEqual(out.stride(), expected.stride())
+
     def test_to_copy_multiple_returns(self):
         class ToCopyReturns(nn.Module):
             def forward(self, x):

From 8ba2a10c282a3add744c148c61ac1855dfbeaa92 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:09:16 -0700
Subject: [PATCH 24/36] fix(lowering): check the attention layout without
 guarding on data-dependent sizes

With decompose_attention=True, torch_tensorrt.dynamo.compile failed when the
attention input had a size known only at runtime, for example after boolean
indexing:

    query = x[mask].transpose(1, 2)
    out = F.scaled_dot_product_attention(query, query, query)

    GuardOnDataDependentSymNode: Could not guard on data-dependent
    expression Eq(160*u0, 0)

The attention output layout pass asked `is_contiguous()` on the recorded
output. For a tensor with a data-dependent size that needs a guard, which
export cannot add.

The pass now uses `is_contiguous_or_false`, which returns False instead of
guarding. The dim order that follows already avoids such guards, so a
contiguous output is still left alone, and a fused layout is still preserved.
Preserving it matters: the same model with
`out.transpose(1, 2).contiguous().view(...)` after attention fails with
"Cannot view a tensor" if the pass skips these nodes.

Test: attention over boolean-indexed rows, compiled once and compared with
eager for 1 and 3 selected rows.
---
 .../preserve_attention_output_layout.py       |  3 ++-
 .../lowering/test_aten_lowering_passes.py     | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py
index 11060ff3167..1fa9cd702a2 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py
@@ -4,6 +4,7 @@
 import logging
 
 import torch
+from torch._prims_common import is_contiguous_or_false
 from torch._subclasses.fake_tensor import FakeTensor
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
@@ -18,7 +19,7 @@ def _preserve_layouts(graph: torch.fx.Graph) -> int:
         val = node.meta.get("val")
         # Under torch.vmap the recorded value is a batched tensor, not a FakeTensor,
         # and those nodes lower without this pass.
-        if not isinstance(val, FakeTensor) or val.is_contiguous():
+        if not isinstance(val, FakeTensor) or is_contiguous_or_false(val):
             continue
         dim_order = val.dim_order()
         if dim_order == tuple(range(val.dim())):
diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
index 9448224d9d8..da69ecba3a8 100644
--- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py
+++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
@@ -484,6 +484,28 @@ def forward(self, query, key, value):
             model(*inputs), trt_module(*inputs), rtol=1e-2, atol=1e-2
         )
 
+    def test_attention_on_data_dependent_rows(self):
+        # Boolean indexing gives attention a batch size known only at runtime.
+        class MaskedAttention(torch.nn.Module):
+            def forward(self, x, mask):
+                query = x[mask].transpose(1, 2)
+                return torch.nn.functional.scaled_dot_product_attention(
+                    query, query, query
+                )
+
+        model = MaskedAttention().cuda()
+        x = torch.randn(3, 16, 4, 32).half().cuda()
+        mask = torch.tensor([True, False, True]).cuda()
+        ep = torch.export.export(model, (x, mask))
+        trt_module = torch_tensorrt.dynamo.compile(
+            ep, [x, mask], min_block_size=1, decompose_attention=True
+        )
+        for mask in ([True, False, False], [True, True, True]):
+            mask = torch.tensor(mask).cuda()
+            torch.testing.assert_close(
+                model(x, mask), trt_module(x, mask), rtol=1e-2, atol=1e-2
+            )
+
 
 class TestRewriteEfficientAttention(TestCase):
     def test_force_causal_efficient_attention(self):

From 7969d82511934d6711bb361afbf57148e5897bc5 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:09:40 -0700
Subject: [PATCH 25/36] test(lowering): cover the attention layout pass inside
 a torch.cond branch

The attention output layout pass also visits torch.cond branches, but only
attention inside a torch.no_grad() region was tested. Without the pass, a
branch that runs attention on a transposed query and then
`out.transpose(1, 2).contiguous().view(...)` fails to compile with
"Cannot view a tensor".

Test: that branch, compiled with decompose_attention=True and compared with
eager for both values of the predicate.
---
 .../lowering/test_aten_lowering_passes.py     | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
index da69ecba3a8..2f4e18e2813 100644
--- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py
+++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
@@ -484,6 +484,35 @@ def forward(self, query, key, value):
             model(*inputs), trt_module(*inputs), rtol=1e-2, atol=1e-2
         )
 
+    def test_view_after_attention_in_cond_branch(self):
+        class CondViewAfterAttention(torch.nn.Module):
+            def forward(self, pred, query, key, value):
+                def attend(query, key, value):
+                    query, key, value = (t.transpose(1, 2) for t in (query, key, value))
+                    out = torch.nn.functional.scaled_dot_product_attention(
+                        query, key, value
+                    )
+                    out = out.transpose(1, 2).contiguous()
+                    return out.view(out.shape[0], out.shape[1], -1)
+
+                def flatten(query, key, value):
+                    return query.flatten(2).clone()
+
+                return torch.cond(pred, attend, flatten, (query, key, value))
+
+        inputs = [torch.randn(2, 16, 4, 32).half().cuda() for _ in range(3)]
+        model = CondViewAfterAttention().cuda()
+        pred = torch.tensor(True).cuda()
+        ep = torch.export.export(model, (pred, *inputs))
+        trt_module = torch_tensorrt.dynamo.compile(
+            ep, [pred, *inputs], min_block_size=1, decompose_attention=True
+        )
+        for pred in (True, False):
+            pred = torch.tensor(pred).cuda()
+            torch.testing.assert_close(
+                model(pred, *inputs), trt_module(pred, *inputs), rtol=1e-2, atol=1e-2
+            )
+
     def test_attention_on_data_dependent_rows(self):
         # Boolean indexing gives attention a batch size known only at runtime.
         class MaskedAttention(torch.nn.Module):

From acd3fcedda62f173c20a26670acd38abb2744827 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:09:41 -0700
Subject: [PATCH 26/36] test(lowering): check that the attention layout test
 sees a fused layout

The attention layout test only fails without the fix when the attention
kernel picked at export returns a layout that is not contiguous. With only
the math kernel enabled, the output is contiguous, export keeps the
`.contiguous()` call, and the test passes without the fix.

The test now checks that the exported attention output is not contiguous, so
it fails loudly instead of passing without testing anything.
---
 tests/py/dynamo/lowering/test_aten_lowering_passes.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
index 2f4e18e2813..110a5e644d0 100644
--- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py
+++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
@@ -477,6 +477,16 @@ def forward(self, query, key, value):
         inputs = [torch.randn(2, 16, 4, 32).half().cuda() for _ in range(3)]
         model = ViewAfterAttention().cuda()
         ep = torch.export.export(model, tuple(inputs))
+        # With a contiguous attention output, export keeps the .contiguous() and
+        # this test passes without the pass.
+        (attention,) = (
+            node
+            for module in ep.graph_module.modules()
+            if isinstance(module, torch.fx.GraphModule)
+            for node in module.graph.nodes
+            if node.target == torch.ops.aten.scaled_dot_product_attention.default
+        )
+        self.assertFalse(attention.meta["val"].is_contiguous())
         trt_module = torch_tensorrt.dynamo.compile(
             ep, inputs, min_block_size=1, decompose_attention=True
         )

From 0b12bb947b923b43e59d122cf4ee10583544f523 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:50:54 -0700
Subject: [PATCH 27/36] fix(executorch): keep the output shape error when
 export refuses an engine

For an engine marked as needing an output allocator, ExecuTorch export asks
the engine for its output shapes and refuses it only if a shape stays
unknown. The C++ engine reports an unknown shape by raising a RuntimeError,
so export treated every RuntimeError from that call as an unknown shape. A
different failure, such as "Unable to create TensorRT execution context",
was dropped: the user saw only "ExecuTorch export does not support TensorRT
engines that require an output allocator", with no sign of the real error.

The engine is still refused, because export cannot tell it is safe, but the
refusal now carries the original error as its cause, so the traceback shows
both.

Test: an engine whose output shape query raises is refused, and the refusal's
cause is that error.
---
 py/torch_tensorrt/executorch/_export_utils.py | 46 ++++++++++++-------
 tests/py/dynamo/executorch/test_export.py     | 16 +++++++
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py
index db4e7c0a111..ca5e0fd112f 100644
--- a/py/torch_tensorrt/executorch/_export_utils.py
+++ b/py/torch_tensorrt/executorch/_export_utils.py
@@ -111,16 +111,20 @@ def _outputs_depend_on_data(engine_obj: Any, node: Any) -> bool:
         if not all(has_guarding_hint(dim) for dim in val.shape):
             return True
         input_shapes.append([guarding_hint_or_throw(dim) for dim in val.shape])
-    try:
-        output_shapes = engine_obj.infer_outputs(input_shapes)
-    except RuntimeError:
-        # The C++ engine cannot create a tensor for an unknown dimension.
-        return True
+    output_shapes = engine_obj.infer_outputs(input_shapes)
     return any(-1 in tuple(getattr(shape, "shape", shape)) for shape in output_shapes)
 
 
-def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> None:
-    """Reject engine configurations unsupported by ExecuTorch export."""
+def validate_engine_info(
+    engine_info: Sequence[Any],
+    *,
+    node_name: str = "",
+    cause: BaseException | None = None,
+) -> None:
+    """Reject engine configurations unsupported by ExecuTorch export.
+
+    ``cause`` is chained to the refusal, so the error behind it stays visible.
+    """
     from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
         REQUIRES_OUTPUT_ALLOCATOR_IDX,
     )
@@ -134,7 +138,7 @@ def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") ->
             "ExecuTorch export does not support TensorRT engines that require "
             "an output allocator (data-dependent output shapes)"
             f"{node_suffix}."
-        )
+        ) from cause
 
 
 def _schema_name(node: Any) -> str:
@@ -294,20 +298,28 @@ def validate_engine_program(
             )
             if engine_node is not None:
                 by_engine_node[engine_node] = engine_info
+        inference_error = None
         if (
             node.target is torch.ops.tensorrt.execute_engine.default
             and len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX
             and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1"
-            and not _outputs_depend_on_data(
-                _resolve_engine_object(exported_program, node), node
-            )
         ):
-            # Every output shape is known, so ExecuTorch needs no output allocator.
-            # This record becomes the ExecuTorch engine's arguments, and the backend
-            # refuses a marked one.
-            engine_info = list(engine_info)
-            engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "0"
-        validate_engine_info(engine_info, node_name=node.name)
+            try:
+                depends_on_data = _outputs_depend_on_data(
+                    _resolve_engine_object(exported_program, node), node
+                )
+            except RuntimeError as error:
+                # The C++ engine raises when it cannot create an output for an
+                # unknown dimension, but also when anything else fails, so the
+                # refusal keeps the error as its cause.
+                depends_on_data, inference_error = True, error
+            if not depends_on_data:
+                # Every output shape is known, so ExecuTorch needs no output
+                # allocator. This record becomes the ExecuTorch engine's
+                # arguments, and the backend refuses a marked one.
+                engine_info = list(engine_info)
+                engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "0"
+        validate_engine_info(engine_info, node_name=node.name, cause=inference_error)
         if resolved is not None:
             resolved[node.name] = engine_info
         count += 1
diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py
index d95e987a586..7cfdeecef57 100644
--- a/tests/py/dynamo/executorch/test_export.py
+++ b/tests/py/dynamo/executorch/test_export.py
@@ -489,6 +489,22 @@ def test_validate_engine_program_rejects_output_allocator_without_mutation():
     assert program.state_dict == {}
 
 
+@pytest.mark.unit
+def test_validate_engine_program_keeps_the_output_shape_error_as_the_cause():
+    export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils")
+    program, _, input_node, engine = _engine_program(
+        lifted=False, requires_output_allocator=True
+    )
+    input_node.meta["val"] = torch.empty(2, 3)
+    error = RuntimeError("Unable to create TensorRT execution context")
+    engine.infer_outputs = MagicMock(side_effect=error)
+
+    with pytest.raises(RuntimeError, match="output allocator") as refusal:
+        export_utils.validate_engine_program(program)
+
+    assert refusal.value.__cause__ is error
+
+
 @pytest.mark.unit
 def test_export_returns_edge_and_forwards_all_options(monkeypatch):
     export_module, lower = _patch_lowering(monkeypatch)

From 19f0c97aa7e61c1c971d4549f72578c7858ec6a8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:51:56 -0700
Subject: [PATCH 28/36] refactor(executorch): name the export output shape
 check for what it proves

The helper that decides whether an engine marked as needing an output
allocator can still be exported was called _outputs_depend_on_data. It
returns True for anything it cannot prove static, such as an input with no
example shape, so the name said more than the code knows. It is now
_output_shapes_are_known, which returns True only when every output shape is
proven, and the export check reads "if the shapes are known, clear the mark".

The validate_engine_program docstring still said validation reads only
metadata. For a marked engine it now also asks the live engine for its output
shapes, which sets input shapes on the engine's execution context and can
create that context. The docstring says so.

No behavior change.
---
 py/torch_tensorrt/executorch/_export_utils.py | 35 +++++++++++--------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py
index ca5e0fd112f..7467f5c6548 100644
--- a/py/torch_tensorrt/executorch/_export_utils.py
+++ b/py/torch_tensorrt/executorch/_export_utils.py
@@ -96,23 +96,27 @@ def get_engine_info_from_state(
     return list(state[0] if isinstance(state, tuple) else state)
 
 
-def _outputs_depend_on_data(engine_obj: Any, node: Any) -> bool:
-    """Whether the engine leaves an output shape unknown for this node's input shapes.
+def _output_shapes_are_known(engine_obj: Any, node: Any) -> bool:
+    """Whether the engine knows every output shape for this node's input shapes.
 
     Torch-TensorRT marks an engine as needing an output allocator when any of its
     operations has a data-dependent shape, which also keeps CUDA graphs off for it.
-    Only an output with such a shape needs the allocator ExecuTorch lacks.
+    Only an output with such a shape needs the allocator ExecuTorch lacks. An input
+    without a concrete example shape makes the answer False, since nothing proves
+    the outputs known.
     """
     input_shapes = []
     for arg in node.args[0]:
         val = arg.meta.get("val") if isinstance(arg, torch.fx.Node) else None
         if not isinstance(val, torch.Tensor):
-            return True
+            return False
         if not all(has_guarding_hint(dim) for dim in val.shape):
-            return True
+            return False
         input_shapes.append([guarding_hint_or_throw(dim) for dim in val.shape])
     output_shapes = engine_obj.infer_outputs(input_shapes)
-    return any(-1 in tuple(getattr(shape, "shape", shape)) for shape in output_shapes)
+    return all(
+        -1 not in tuple(getattr(shape, "shape", shape)) for shape in output_shapes
+    )
 
 
 def validate_engine_info(
@@ -262,8 +266,11 @@ def validate_engine_program(
     Resolving an engine serializes it, so pass a dict as ``resolved`` to keep what was
     resolved here and let the rewrite reuse it instead of serializing a second time.
 
-    Validation reads only metadata, so the records kept in ``resolved`` are
-    metadata-only: their ENGINE_IDX slot is empty. The rewrite reuses them for the
+    Validation reads each engine's metadata. For an engine marked as needing an
+    output allocator it also asks the live engine for its output shapes, which sets
+    the node's input shapes on the engine's execution context and creates the context
+    if no run has yet. The records kept in ``resolved`` are metadata-only: their
+    ENGINE_IDX slot is empty. The rewrite reuses them for the
     plain-string slots and takes the engine itself from :func:`_resolve_engine_tensor`,
     which never reads a cached record's ENGINE_IDX -- so a metadata-only record is
     never cross-served as the engine bytes.
@@ -305,18 +312,18 @@ def validate_engine_program(
             and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1"
         ):
             try:
-                depends_on_data = _outputs_depend_on_data(
+                shapes_known = _output_shapes_are_known(
                     _resolve_engine_object(exported_program, node), node
                 )
             except RuntimeError as error:
                 # The C++ engine raises when it cannot create an output for an
                 # unknown dimension, but also when anything else fails, so the
                 # refusal keeps the error as its cause.
-                depends_on_data, inference_error = True, error
-            if not depends_on_data:
-                # Every output shape is known, so ExecuTorch needs no output
-                # allocator. This record becomes the ExecuTorch engine's
-                # arguments, and the backend refuses a marked one.
+                shapes_known, inference_error = False, error
+            if shapes_known:
+                # ExecuTorch needs no output allocator then. This record becomes
+                # the ExecuTorch engine's arguments, and the backend refuses a
+                # marked one.
                 engine_info = list(engine_info)
                 engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "0"
         validate_engine_info(engine_info, node_name=node.name, cause=inference_error)

From 44042cd62558593a36f763e49f7ff07018008637 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:51:57 -0700
Subject: [PATCH 29/36] test(executorch): mark the data-dependent shape export
 test as a unit test

Every other test in test_export.py carries @pytest.mark.unit, so a run with
-m unit skipped only this one.
---
 tests/py/dynamo/executorch/test_export.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py
index 7cfdeecef57..2d26ef22675 100644
--- a/tests/py/dynamo/executorch/test_export.py
+++ b/tests/py/dynamo/executorch/test_export.py
@@ -873,6 +873,7 @@ def forward(self, value):
     assert output.stat().st_size > 0
 
 
+@pytest.mark.unit
 def test_engine_with_data_dependent_shape_only_inside_exports_a_pte(tmp_path):
     """The boolean mask selection has a data-dependent shape, the engine output does not."""
 

From 205218aa49ed646bdcc37000aff986b8dd030151 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:53:47 -0700
Subject: [PATCH 30/36] test(cond): check the values of a conditional whose
 lowered branch is installed

test_cond_keeps_branches_returned_by_lowering checked only that the lowered
conditional is convertible. It now also runs the lowered module for both
predicates and compares the result with the original model, so a branch that
is installed but computes the wrong values fails the test.
---
 tests/py/dynamo/models/test_cond.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/py/dynamo/models/test_cond.py b/tests/py/dynamo/models/test_cond.py
index aeedb561be3..d60eb477edf 100644
--- a/tests/py/dynamo/models/test_cond.py
+++ b/tests/py/dynamo/models/test_cond.py
@@ -141,9 +141,11 @@ def test_cond_is_supported_when_branches_are_convertible(self):
     def test_cond_keeps_branches_returned_by_lowering(self):
         # The complex lowering returns a new branch module rather than editing it.
         x = torch.randn(2, 4)
-        pred = torch.tensor(True)
-        gm = _export_lowered(_ComplexCond(), (x, pred))
+        model = _ComplexCond()
+        gm = _export_lowered(model, (x, torch.tensor(True)))
         self.assertTrue(_cond_node(gm) in CONVERTERS)
+        for pred in (torch.tensor(True), torch.tensor(False)):
+            torch.testing.assert_close(gm(x, pred), model(x, pred))
 
     def test_cond_keeps_branch_when_lowering_changes_output_layout(self):
         pred = torch.tensor(True)

From 3bc88c066716c8d89e02a900fcc6f0e0562f3bab Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 15:54:13 -0700
Subject: [PATCH 31/36] fix(lowering): compare sizes and device before
 installing a lowered cond branch

post_lowering installs the module that lowering returns for a torch.cond
branch only if its inputs and outputs keep their dtype and memory layout,
because the cond node was traced against the old branch. The check compared
dtype, rank, strides, storage offset and storage size, but not the sizes or
the device. So a branch that came back with a (2, 4) output where the
original returned (3, 4) from the same storage, or on another device, would
have been installed, and the cond node would describe the wrong tensor. No
current pass returns such a branch, so this keeps the check complete rather
than fixing a seen failure.

The check now also requires equal sizes, proven the same way as the strides,
and the same device.

Test: the branch check rejects two modules whose tensors differ only in sizes,
and two that differ only in device.
---
 .../lowering/passes/_aten_lowering_pass.py    |  8 +++++++-
 tests/py/dynamo/models/test_cond.py           | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
index a9733825f3a..132b2c191fd 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
@@ -224,7 +224,11 @@ def value(node: Any) -> Any:
 
 
 def _same_io(a: torch.fx.GraphModule, b: torch.fx.GraphModule) -> bool:
-    """Whether two modules take and return tensors of the same dtype and layout."""
+    """Whether two modules take and return tensors of the same kind and layout.
+
+    Tensors match when their dtype, device, sizes, strides, storage offset and
+    storage size match.
+    """
 
     def same(x: Any, y: Any) -> bool:
         if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor):
@@ -232,7 +236,9 @@ def same(x: Any, y: Any) -> bool:
         # A symbolic value that cannot be proven equal counts as different.
         return (
             x.dtype == y.dtype
+            and x.device == y.device
             and x.dim() == y.dim()
+            and all(statically_known_true(p == q) for p, q in zip(x.shape, y.shape))
             and statically_known_true(
                 x.untyped_storage().nbytes() == y.untyped_storage().nbytes()
             )
diff --git a/tests/py/dynamo/models/test_cond.py b/tests/py/dynamo/models/test_cond.py
index d60eb477edf..4fa93d21a02 100644
--- a/tests/py/dynamo/models/test_cond.py
+++ b/tests/py/dynamo/models/test_cond.py
@@ -12,6 +12,7 @@
     post_lowering,
     pre_export_lowering,
 )
+from torch_tensorrt.dynamo.lowering.passes._aten_lowering_pass import _same_io
 
 
 def _cond_node(gm: torch.fx.GraphModule) -> torch.fx.Node:
@@ -174,6 +175,25 @@ def test_cond_keeps_branch_when_lowering_changes_output_dtype(self):
         gm = _export_lowered(model, (x, pred))
         torch.testing.assert_close(gm(x, pred), model(x, pred))
 
+    def test_cond_branch_check_compares_sizes_and_device(self):
+        def identity(value):
+            graph = torch.fx.Graph()
+            node = graph.placeholder("x")
+            node.meta["val"] = value
+            graph.output(node)
+            return torch.fx.GraphModule(torch.nn.Module(), graph)
+
+        storage = torch.empty(12)
+        cases = {
+            # Same dtype, strides, storage size and offset.
+            "sizes": (storage[:8].view(2, 4), storage.view(3, 4)),
+            "device": (torch.empty(2, 4), torch.empty(2, 4, device="meta")),
+        }
+        for name, (a, b) in cases.items():
+            with self.subTest(name):
+                self.assertTrue(_same_io(identity(a), identity(a)))
+                self.assertFalse(_same_io(identity(a), identity(b)))
+
     def test_cond_falls_back_when_branch_has_unsupported_op(self):
         x = torch.randn(1, 4).abs() + 0.1
         pred = torch.tensor(True)

From 750364489aa30a12355df42d428f8dbcf6e68496 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:00:47 -0700
Subject: [PATCH 32/36] fix(lowering): keep scatter_add positions live when
 constant folding keeps src

The scatter_add_positions rule keeps the positions arange of the integer
scatter_add decomposition out of constant folding, but only when src varies
at run time. Otherwise the whole update folds to an output-sized constant.
It decided that src varies only when src reads an input or a random op. Two
more kinds of src stay in the graph after folding:

- a node that another rule keeps out of folding, such as an arange shared by
  a causal attention mask and src while attention_mask_arange is enabled;
- a view whose storage the graph returns or mutates, which the folder keeps
  to preserve the alias.

For those the rule let the positions fold while src stayed live, so the
comparison of every index with every position became an engine constant.
With arange(96) as both a causal mask and src, and a constant index into
8192 outputs, that constant took 6.3 MB of a 6.3 MB engine; with the
positions kept, the engine is about 50 KB.

The rule now treats both kinds as varying at run time. The exclusion pass
applies one rule to the whole graph before the next, so the scatter_add rule
sees the attention marks wherever they sit in the graph.

Test: two new cases of the scatter_add positions test, src from an attention
mask arange and src returned as a view, keep the positions arange.
---
 .../constant_fold_exclusions/scatter_add.py   | 18 +++++++++++--
 .../passes/mark_constant_fold_exclusions.py   | 10 ++++---
 .../py/dynamo/lowering/test_decompositions.py | 26 ++++++++++++++++---
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
index a9fc83af230..65f04e7f425 100644
--- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
@@ -7,7 +7,10 @@
 from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor
 from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot
 
-from ._core import register_constant_fold_exclusion_rule
+from ._core import (
+    CONSTANT_FOLD_EXCLUSION_META_KEY,
+    register_constant_fold_exclusion_rule,
+)
 
 SCATTER_ADD_POSITIONS_RULE_ID = "scatter_add_positions"
 _POSITIONS_TAG = "_torch_tensorrt_scatter_add_positions"
@@ -35,7 +38,16 @@ def exclude_scatter_add_positions_from_constant_fold(positions: torch.Tensor) ->
 
 
 def _varies_at_runtime(node: torch.fx.Node) -> bool:
-    """Whether constant folding cannot compute this node ahead of time."""
+    """Whether this node reads a value that constant folding leaves in the graph.
+
+    That is an input, a random op, a node an earlier rule keeps out of folding, or
+    a view whose storage the graph returns or mutates.
+    """
+    from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import (
+        _aliases_input,
+        _storage_escapes,
+    )
+
     nodes_to_visit, visited = [node], set()
     while nodes_to_visit:
         current = nodes_to_visit.pop()
@@ -43,6 +55,8 @@ def _varies_at_runtime(node: torch.fx.Node) -> bool:
         if (
             current.op == "placeholder"
             or torch.Tag.nondeterministic_seeded in target_tags
+            or current.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY)
+            or (_aliases_input(current) and _storage_escapes(current))
         ):
             return True
         if current not in visited:
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
index 311d23fe5e9..7feb16fd4e5 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py
@@ -28,10 +28,12 @@ def mark_constant_fold_exclusions(
         settings.disabled_constant_fold_exclusions if settings is not None else ()
     )
 
-    for node in gm.graph.nodes:
-        for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items():
-            if rule_id in disabled_rule_ids:
-                continue
+    # One rule covers the whole graph before the next starts, so a rule can read
+    # the marks of the rules registered before it.
+    for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items():
+        if rule_id in disabled_rule_ids:
+            continue
+        for node in gm.graph.nodes:
             _mark_constant_fold_exclusion(rule(node), rule_id)
 
     if disabled_rule_ids:
diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py
index 573a6f2939a..df3fbde79f7 100644
--- a/tests/py/dynamo/lowering/test_decompositions.py
+++ b/tests/py/dynamo/lowering/test_decompositions.py
@@ -1413,6 +1413,9 @@ def forward(self, base, index, src):
             ("decomposed_twice", (), 2, "input", True),
             ("constant_src_folds", (), 1, "ones", False),
             ("random_src", (), 1, "randint", True),
+            # Constant src that constant folding still leaves in the graph.
+            ("attention_mask_src", (), 1, "attention_mask_arange", True),
+            ("returned_src_view", (), 1, "returned_view", True),
         ]
     )
     def test_scatter_add_constant_index_positions_not_folded(
@@ -1433,7 +1436,23 @@ def forward(self, input, src):
                     src = torch.ones_like(self.index)
                 elif src_kind == "randint":
                     src = torch.randint_like(self.index, 5)
-                return input.scatter_add(0, self.index, src)
+                elif src_kind == "attention_mask_arange":
+                    sequence = torch.arange(8, device=input.device)
+                    src = sequence.expand_as(self.index)
+                elif src_kind == "returned_view":
+                    src = self.index.t().t()
+                result = input.scatter_add(0, self.index, src)
+                if src_kind == "attention_mask_arange":
+                    # The attention comes after the scatter_add positions in the graph.
+                    mask = sequence.unsqueeze(0) <= sequence.unsqueeze(1)
+                    query = input[:8].float().view(1, 1, 8, 8)
+                    attention = torch.ops.aten.scaled_dot_product_attention.default(
+                        query, query, query, mask
+                    )
+                    return result, attention
+                if src_kind == "returned_view":
+                    return result, src
+                return result
 
         exported = torch.export.export(
             ConstantIndexScatterAdd().cuda(),
@@ -1450,11 +1469,12 @@ def forward(self, input, src):
                 disabled_constant_fold_exclusions=disabled_constant_fold_exclusions
             ),
         )
-        has_arange = any(
+        has_positions = any(
             getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange
+            and node.meta["val"].shape == (64,)
             for node in gm.graph.nodes
         )
-        self.assertEqual(has_arange, positions_kept)
+        self.assertEqual(has_positions, positions_kept)
 
     def test_scatter_add_large_comparison_keeps_loop(self):
         class ScatterAdd(torch.nn.Module):

From e1abd894b7e47d74fc193130cd2abc2c5876c7da Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:05:21 -0700
Subject: [PATCH 33/36] docs(lowering): clarify two scatter_add comments

Two comments in the integer scatter_add change were easy to misread:

- "It specializes a dynamic extent" read as if this branch handled dynamic
  extents. It now says the loop would specialize one, which is why
  filter_decomposition_table keeps such an op out of decomposition.
- The positions tag docstring said folding stores "the whole comparison" as a
  constant, which read as if the multiply by src folded too. It now names the
  0/1 mask that src is multiplied by, which folds even when src varies.

No behavior change.
---
 py/torch_tensorrt/dynamo/lowering/_decompositions.py         | 4 +++-
 .../dynamo/lowering/constant_fold_exclusions/scatter_add.py  | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
index c75fa9e998e..1bfcfa5d2c4 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
@@ -349,7 +349,9 @@ def _use_integer_scatter_add(
     elif index.dim() and isinstance(index.shape[dim], int):
         extent = input_tensor.shape[dim]
     else:
-        # Otherwise keep the loop, as before. It specializes a dynamic extent.
+        # Otherwise keep the loop, as before. On a dynamic extent the loop would
+        # specialize it, which is why filter_decomposition_table keeps such an op
+        # out of decomposition on the compile paths.
         return False
     comparisons = index.numel() * extent
     return statically_known_true(comparisons <= _MAX_INTEGER_SCATTER_ADD_COMPARISONS)
diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
index 65f04e7f425..ab6e902fef8 100644
--- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
+++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py
@@ -19,8 +19,9 @@
 def exclude_scatter_add_positions_from_constant_fold(positions: torch.Tensor) -> None:
     """Tag the positions arange of the scatter_add decomposition.
 
-    With a constant index, folding it would store the whole comparison of every
-    index with every position as an engine constant. Retraces keep
+    With a constant index, folding it would store the comparison of every index
+    with every position, the 0/1 mask that src is multiplied by, as an engine
+    constant, even when src varies at run time. Retraces keep
     ``node.meta["custom"]``, so the rule below finds the tag after the
     decomposed op is gone, for example after the complex decomposition retrace
     or a second ``run_decompositions``.

From 7a9e10b11ed4867295784158a7d4dedbd78da11f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:13:16 -0700
Subject: [PATCH 34/36] fix(lowering): free offloaded constant copies before
 the legacy complex pass

With offload_module_to_cpu, the complex decomposition copies each folded
constant that lives on the CPU to its device for the retrace. When the
retrace fails, the adapter falls back to the legacy complex pass. It did so
inside the except block, where the exception still keeps the failed
retrace's frames alive, and those frames hold the device copies. So the
legacy pass, which can run for a long time on a large model, ran with a
second device copy of every offloaded constant. For a 256 MiB folded
constant and a complex log2, which the retrace does not support, that was
256 MiB more GPU memory while the legacy pass ran.

The fallback now runs after the except block, when the frames and the copies
are gone.

Test: a retrace that fails while its frame holds the constants falls back to
the legacy pass with no more GPU memory allocated than before the adapter ran.
---
 .../passes/complex_decomposition_adapter.py   | 10 ++--
 .../dynamo/lowering/test_constant_folding.py  | 55 +++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
index aab92106ef0..d598dcce680 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py
@@ -30,7 +30,7 @@
 import gc
 import logging
 from contextlib import contextmanager
-from typing import Any, Iterator
+from typing import Any, Iterator, Optional
 
 import torch
 from torch._export.utils import _detect_fake_mode_from_gm
@@ -110,6 +110,7 @@ def complex_decomposition_adapter(
     # codegen then unflattens again.
     gm.graph._codegen = torch.fx.graph.CodeGen()
     gm.recompile()
+    decomposed_gm: Optional[GraphModule] = None
     try:
         # Both retraces return a new graph. Keep the original for the existing
         # unsupported-op fallback rather than replacing it before success.
@@ -136,11 +137,12 @@ def complex_decomposition_adapter(
             f"complex decomposition failed ({type(e).__name__}: {e}); "
             "falling back to the legacy complex_graph_detection pass."
         )
-        gm.graph._codegen = original_codegen
-        gm.recompile()
-        return complex_graph_detection(gm, settings)
     gm.graph._codegen = original_codegen  # the caller's module, as it was
     gm.recompile()
+    if decomposed_gm is None:
+        # Not inside the except block: there the failed retrace's frames are still
+        # alive, and with them its device copies of the offloaded constants.
+        return complex_graph_detection(gm, settings)
     _restore_offloaded_constants(decomposed_gm, offloaded)
     if offloaded:
         # The retrace leaves the device copies in reference cycles. Free them now
diff --git a/tests/py/dynamo/lowering/test_constant_folding.py b/tests/py/dynamo/lowering/test_constant_folding.py
index 122e58f3c50..6df25a789a0 100644
--- a/tests/py/dynamo/lowering/test_constant_folding.py
+++ b/tests/py/dynamo/lowering/test_constant_folding.py
@@ -202,6 +202,61 @@ def forward(self, x):
         torch.testing.assert_close(actual, expected)
         self.assertEqual(model.scale.device.type, "cpu")
 
+    @unittest.skipIf(
+        not torch.cuda.is_available()
+        or not complex_decomposition_adapter.has_complex_decomposition(),
+        "needs CUDA and the upstream complex decomposition",
+    )
+    def test_complex_fallback_does_not_hold_device_copies_of_offloaded_constants(
+        self,
+    ):
+        from torch._functorch._aot_autograd import complex_decomposition
+
+        class Rotary(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                angles = torch.outer(torch.arange(256.0), torch.ones(512))
+                self.register_buffer(
+                    "freqs", torch.polar(torch.ones_like(angles), angles)
+                )
+                self.scale = torch.nn.Parameter(torch.ones(512))
+
+            def forward(self, x):
+                z = torch.view_as_complex(x.reshape(256, 512, 2))
+                return torch.view_as_real(z * (self.freqs * self.scale.mean()))
+
+        settings = CompilationSettings(offload_module_to_cpu=True)
+        example_input = torch.randn(256, 1024, device="cuda")
+        exported_program = torch.export.export(Rotary().cuda(), (example_input,))
+        # Folds freqs * scale into a 1 MiB constant kept on the CPU.
+        gm = constant_fold(exported_program.module(), settings)
+
+        def failing_retrace(gm, flat_args, decompositions):
+            # Like a real retrace, its frame holds the constants it was handed.
+            constants = [
+                getattr(gm, n.target) for n in gm.graph.nodes if n.op == "get_attr"
+            ]
+            raise NotImplementedError(f"unsupported op, {len(constants)} constants")
+
+        allocated = {}
+
+        def legacy_pass(gm, settings):
+            allocated["fallback"] = torch.cuda.memory_allocated()
+            return gm
+
+        with (
+            mock.patch.object(
+                complex_decomposition, "decompose_complex_in_graph", failing_retrace
+            ),
+            mock.patch.object(
+                complex_decomposition_adapter, "complex_graph_detection", legacy_pass
+            ),
+        ):
+            allocated["before"] = torch.cuda.memory_allocated()
+            complex_decomposition_adapter.complex_decomposition_adapter(gm, settings)
+
+        self.assertEqual(allocated["fallback"], allocated["before"])
+
 
 if __name__ == "__main__":
     run_tests()

From a7a0e4b605645eb5b9739c7ea8b0b2fe5c051afc Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 17:55:30 -0700
Subject: [PATCH 35/36] test(conversion): accept the IndexError PyTorch raises
 for a mask count mismatch

The data-dependent boolean mask test checks that a mask whose True count
differs from the value rows fails, as it does in PyTorch. It expected a
RuntimeError. When the index_put runs in PyTorch instead of TensorRT, as it
does on TensorRT-RTX for the SetRowsAtColumns case, PyTorch raises IndexError
("indexing tensors could not be broadcast together"), so the test failed
although the behaviour matched PyTorch.

The test now accepts either error.
---
 tests/py/dynamo/conversion/test_index_put_aten.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py
index 02f38a7c2f9..5864edb8905 100644
--- a/tests/py/dynamo/conversion/test_index_put_aten.py
+++ b/tests/py/dynamo/conversion/test_index_put_aten.py
@@ -567,10 +567,11 @@ def forward(self, data, mask, values):
                     torch.testing.assert_close(
                         compiled(data, run_mask, values), model(data, run_mask, values)
                     )
-                # A True count that differs from the value rows fails, as in PyTorch.
+                # A True count that differs from the value rows fails, as in PyTorch,
+                # which raises IndexError when the op runs outside TensorRT.
                 fewer = case_mask.clone()
                 fewer[..., 6] = False
-                with self.assertRaises(RuntimeError):
+                with self.assertRaises((RuntimeError, IndexError)):
                     compiled(data, fewer, values)
 
     def test_updates_are_cast_to_int64_destination_dtype(self):

From 8808d98045146659b83d69cc8b7f911aaefc76f0 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 18:15:49 -0700
Subject: [PATCH 36/36] test(executorch): pin the error code and the version
 source in the load error tests

The no-GPU test checked only that loading failed with the new message. It
now also checks the NotFound error (0x20) that an unavailable backend gives.

The version test exported and loaded with the same TensorRT, so it could not
tell whether the "exported with" half of the message came from the program or
from the running TensorRT. It now writes a stand-in export version of the same
length into the program and checks that the message names it next to the
running version.
---
 tests/py/dynamo/executorch/test_load_errors.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/py/dynamo/executorch/test_load_errors.py b/tests/py/dynamo/executorch/test_load_errors.py
index 71cfc9f9cb5..3e19982536f 100644
--- a/tests/py/dynamo/executorch/test_load_errors.py
+++ b/tests/py/dynamo/executorch/test_load_errors.py
@@ -76,7 +76,10 @@ def test_loading_with_no_visible_gpu_fails_with_a_clear_error(tmp_path):
         timeout=600,
     )
     assert result.returncode == 3, (result.returncode, result.stderr[-2000:])
-    assert "CUDA is unavailable" in result.stdout + result.stderr
+    output = result.stdout + result.stderr
+    assert "CUDA is unavailable" in output
+    # The backend reports itself unavailable, which ExecuTorch turns into NotFound (0x20).
+    assert "0x20" in output or "0x:20" in output, output[-2000:]
 
 
 def test_an_engine_tensorrt_cannot_read_names_both_tensorrt_versions(tmp_path, capfd):
@@ -98,6 +101,12 @@ def test_an_engine_tensorrt_cannot_read_names_both_tensorrt_versions(tmp_path, c
     start = data.find(engine)
     assert start > 0 and data.find(engine, start + 1) == -1
     data[start : start + 64] = bytes(64)
+    # A stand-in export version of the same length, so the log can only have read it from the
+    # program and not from the TensorRT that is loading it.
+    exported = tensorrt.__version__.encode()
+    stand_in = b".".join(b"9" * len(part) for part in exported.split(b"."))
+    assert data.count(exported) == 1
+    data[data.find(exported) : data.find(exported) + len(exported)] = stand_in
     pte.write_bytes(data)
 
     program = Runtime.get().load_program(pte)
@@ -105,5 +114,5 @@ def test_an_engine_tensorrt_cannot_read_names_both_tensorrt_versions(tmp_path, c
         program.load_method("forward")
     log = str(raised.value) + capfd.readouterr().err
     runtime_version = ".".join(tensorrt.__version__.split(".")[:3])
-    assert f"exported with TensorRT {tensorrt.__version__}" in log, log
+    assert f"exported with TensorRT {stand_in.decode()}" in log, log
     assert f"this runtime is TensorRT {runtime_version}" in log, log