From 0e9ea57f5c4eb84f761098c2a60c77bcec49c271 Mon Sep 17 00:00:00 2001
From: Michael Williams <mww565@icloud.com>
Date: Mon, 17 Aug 2026 12:39:22 -0700
Subject: [PATCH 1/5] Merged main

---
 .../lowering/passes/constant_folding.py       | 42 ++++++++
 .../test_constant_folding_mutable_outputs.py  | 97 +++++++++++++++++++
 2 files changed, 139 insertions(+)
 create mode 100644 tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
index 60dec56f3b..6bac32fb1a 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
@@ -58,6 +58,8 @@ def constant_fold(
         gm.graph.erase_node(node)
 
     gm = clean_up_graph_after_modifications(gm)
+    gm = _clone_folded_constants_at_outputs(gm)
+
     # Delete the constant folder instance which holds GPU memory
     del cf
 
@@ -99,6 +101,46 @@ def replace_node_with_constant(
     setattr(gm, qualname, constant)
 
 
+def _clone_folded_constants_at_outputs(
+    gm: torch.fx.GraphModule,
+) -> torch.fx.GraphModule:
+    """If a folded buffer is returned from the graph, return a clone instead.
+
+    Eager code after a graph break may mutate outputs in-place. Without a clone,
+    that mutates self._frozen_param* and poisons later calls.
+    """
+    output_node = next(n for n in gm.graph.nodes if n.op == "output")
+    # FX output args are typically a tuple/list of returned values
+    out_args = output_node.args[0]
+    if not isinstance(out_args, (tuple, list)):
+        out_args = (out_args,)
+
+    new_outs = []
+    changed = False
+    for out in out_args:
+        if (
+            isinstance(out, torch.fx.Node)
+            and out.op == "get_attr"
+            and str(out.target).startswith("_frozen_param")
+        ):
+            with gm.graph.inserting_before(output_node):
+                cloned = gm.graph.call_function(torch.clone, args=(out,))
+                # keep meta if present
+                if hasattr(out, "meta"):
+                    cloned.meta.update(out.meta)
+            new_outs.append(cloned)
+            changed = True
+        else:
+            new_outs.append(out)
+
+    if changed:
+        output_node.args = (tuple(new_outs),)
+        gm.graph.lint()
+        gm.recompile()
+
+    return gm
+
+
 # TODO: Delete this class when the following code is fixed in nightly:
 # https://github.com/pytorch/pytorch/blob/4b881b0da390c1290bb12850ef9daad6f6eb2cb6/torch/_inductor/constant_folding.py#L53-L63
 class _TorchTensorRTConstantFolder(ConstantFolder):  # type: ignore[misc]
diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
new file mode 100644
index 0000000000..d538b3703d
--- /dev/null
+++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
@@ -0,0 +1,97 @@
+# Regression tests for https://github.com/pytorch/TensorRT/issues/4466
+#
+# Constant folding may replace input-independent factories (e.g. torch.zeros)
+# with persistent _frozen_param* attributes. If those are returned across a
+# graph break and mutated in eager code, later calls must not observe that
+# mutation.
+
+import torch
+import torch_tensorrt  # noqa: F401  # Registers the "tensorrt" backend
+from torch.fx import Graph, GraphModule
+from torch.testing._internal.common_utils import TestCase, run_tests
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold
+
+
+@torch._dynamo.disable
+def mutate_accumulators(values, weight, sample):
+    # Runs eagerly after the graph break.
+    values += sample
+    weight += 1
+    return values / weight
+
+
+def accumulate_with_fresh_state(sample):
+    # Input-independent factories; constant folding replaces these with
+    # persistent _frozen_param* attributes.
+    values = torch.zeros((8,), dtype=torch.float32, device="cuda")
+    weight = torch.zeros((8,), dtype=torch.float32, device="cuda")
+    return mutate_accumulators(values, weight, sample)
+
+
+class TestFoldedConstantGraphBreak(TestCase):
+    def tearDown(self):
+        torch._dynamo.reset()
+
+    def test_state_is_fresh_across_calls(self):
+        """End-to-end repro from issue #4466."""
+        compiled = torch.compile(
+            accumulate_with_fresh_state,
+            backend="tensorrt",
+            dynamic=False,
+            options={"min_block_size": 1},
+        )
+
+        first = compiled(torch.ones(8, device="cuda"))
+        second = compiled(torch.full((8,), 3.0, device="cuda"))
+
+        # Eager semantics: each call starts from fresh zeros.
+        # Before the fix, second would be 2.0 from reused mutated state.
+        torch.testing.assert_close(first, torch.ones_like(first))
+        torch.testing.assert_close(second, torch.full_like(second, 3.0))
+
+    def test_constant_fold_clones_frozen_outputs(self):
+        """Folded constants returned as graph outputs must be cloned."""
+        g = Graph()
+        with g.inserting_after():
+            values = g.call_function(
+                torch.ops.aten.zeros.default,
+                args=((8,),),
+                kwargs={"dtype": torch.float32},
+            )
+            weight = g.call_function(
+                torch.ops.aten.zeros.default,
+                args=((8,),),
+                kwargs={"dtype": torch.float32},
+            )
+            g.output((values, weight))
+
+        gm = GraphModule(torch.nn.Module(), g)
+        gm = constant_fold(gm, CompilationSettings())
+
+        # Outputs should be clones of _frozen_param*, not the attrs themselves.
+        output_node = next(n for n in gm.graph.nodes if n.op == "output")
+        outs = output_node.args[0]
+        self.assertEqual(len(outs), 2)
+        for out in outs:
+            self.assertEqual(out.op, "call_function")
+            self.assertIn(out.target, (torch.clone, torch.ops.aten.clone.default))
+
+        out0, out1 = gm()
+        frozen = [
+            getattr(gm, name)
+            for name in dir(gm)
+            if name.startswith("_frozen_param")
+            and isinstance(getattr(gm, name), torch.Tensor)
+        ]
+        self.assertGreaterEqual(len(frozen), 2)
+
+        # Mutating returned tensors must not change stored folded constants.
+        out0.add_(1)
+        out1.add_(1)
+        for t in frozen:
+            self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8)))
+
+
+if __name__ == "__main__":
+    run_tests()

From 7287d0651b237661e911a50efad18ea349be439a Mon Sep 17 00:00:00 2001
From: Michael Williams <mww565@icloud.com>
Date: Mon, 17 Aug 2026 13:16:38 -0700
Subject: [PATCH 2/5] Added cache for clones so each folded get_attr is cloned
 once

---
 .../lowering/passes/constant_folding.py       | 15 ++++----
 .../test_constant_folding_mutable_outputs.py  | 36 +++++++++++++++++++
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
index 6bac32fb1a..829d995d55 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
@@ -115,6 +115,7 @@ def _clone_folded_constants_at_outputs(
     if not isinstance(out_args, (tuple, list)):
         out_args = (out_args,)
 
+    clone_cache = {}
     new_outs = []
     changed = False
     for out in out_args:
@@ -123,12 +124,14 @@ def _clone_folded_constants_at_outputs(
             and out.op == "get_attr"
             and str(out.target).startswith("_frozen_param")
         ):
-            with gm.graph.inserting_before(output_node):
-                cloned = gm.graph.call_function(torch.clone, args=(out,))
-                # keep meta if present
-                if hasattr(out, "meta"):
-                    cloned.meta.update(out.meta)
-            new_outs.append(cloned)
+            if out not in clone_cache:
+                with gm.graph.inserting_before(output_node):
+                    cloned = gm.graph.call_function(torch.clone, args=(out,))
+                    # keep meta if present
+                    if hasattr(out, "meta"):
+                        cloned.meta.update(out.meta)
+                clone_cache[out] = cloned
+            new_outs.append(clone_cache[out])
             changed = True
         else:
             new_outs.append(out)
diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
index d538b3703d..b108eb5fa5 100644
--- a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
+++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
@@ -92,6 +92,42 @@ def test_constant_fold_clones_frozen_outputs(self):
         for t in frozen:
             self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8)))
 
+    def test_repeated_folded_outputs_preserve_alias(self):
+        """Same folded value returned twice must still share storage (eager semantics)."""
+        g = Graph()
+        with g.inserting_after():
+            values = g.call_function(
+                torch.ops.aten.zeros.default,
+                args=((8,),),
+                kwargs={"dtype": torch.float32},
+            )
+            g.output((values, values))  # same value twice
+
+        gm = GraphModule(torch.nn.Module(), g)
+        gm = constant_fold(gm, CompilationSettings())
+
+        output_node = next(n for n in gm.graph.nodes if n.op == "output")
+        out0, out1 = output_node.args[0]
+        # One clone node reused for both outputs.
+        self.assertIs(out0, out1)
+        self.assertEqual(out0.op, "call_function")
+        self.assertIn(out0.target, (torch.clone, torch.ops.aten.clone.default))
+
+        a, b = gm()
+        self.assertEqual(a.data_ptr(), b.data_ptr())
+        a.add_(1)
+        self.assertTrue(torch.equal(a, b))  # mutation shared across aliases
+
+        frozen = [
+            getattr(gm, name)
+            for name in dir(gm)
+            if name.startswith("_frozen_param")
+            and isinstance(getattr(gm, name), torch.Tensor)
+        ]
+        self.assertGreaterEqual(len(frozen), 1)
+        for t in frozen:
+            self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8)))
+
 
 if __name__ == "__main__":
     run_tests()

From d9de3884a0d929f27854623398d9e3cb7b35bf49 Mon Sep 17 00:00:00 2001
From: Michael Williams <mww565@icloud.com>
Date: Thu, 20 Aug 2026 16:41:14 -0700
Subject: [PATCH 3/5] Reset folded constructors at graph and partition
 boundaries.

Clone compiler-owned _frozen_param outputs late so in-place mutation cannot poison later calls, while leaving user-owned placeholders aliased.
---
 py/torch_tensorrt/dynamo/_compiler.py         | 10 +++-
 .../lowering/passes/_aten_lowering_pass.py    |  2 +
 .../lowering/passes/constant_folding.py       | 44 ----------------
 .../passes/reset_folded_constructors.py       | 50 +++++++++++++++++++
 .../test_constant_folding_mutable_outputs.py  | 45 ++++++++++++++++-
 5 files changed, 105 insertions(+), 46 deletions(-)
 create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index 3ebac2a21f..2b5695a9fd 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -5,7 +5,6 @@
 import os
 import platform
 import warnings
-
 from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union
 
 import sympy
@@ -48,6 +47,9 @@
     inline_lifted_buffers_into_gm,
     lift_mutated_buffers,
 )
+from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import (
+    reset_folded_constructors,
+)
 from torch_tensorrt.dynamo.partitioning._resource_partitioner import (
     resource_partition,
 )
@@ -1353,6 +1355,12 @@ def preserve_module_specs(
                 f"node_name: {name} does not exist in the submodule node dictionary"
             )
 
+        # Partitioning can expose an internal folded constructor as a new TRT
+        # subgraph output. Give that compiler-owned value fresh storage on each
+        # invocation before downstream eager code can mutate it.
+        submodule = reset_folded_constructors(submodule, settings)
+        setattr(partitioned_module, name, submodule)
+
         # set the submodule metadata back to the parent trt_module_node
         metadata_list = get_output_metadata(submodule)
         assert len(metadata_list) > 0
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 b92165f83f..d008d7d264 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py
@@ -26,6 +26,7 @@
 from .repair_input_as_output import repair_input_as_output
 from .replace_fused_rms_norm import replace_fused_rms_norm
 from .replace_max_pool_with_indices import replace_max_pool_with_indices
+from .reset_folded_constructors import reset_folded_constructors
 from .rule_based_autocast import rule_based_autocast
 
 pre_lowering_pass_list = [
@@ -39,6 +40,7 @@
     remove_input_alias_fixing_clones,
     constant_fold,
     repair_input_as_output,
+    reset_folded_constructors,
     fuse_prims_broadcast,
     replace_max_pool_with_indices,
     remove_assert_nodes,
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
index 829d995d55..fd4b350825 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
@@ -58,7 +58,6 @@ def constant_fold(
         gm.graph.erase_node(node)
 
     gm = clean_up_graph_after_modifications(gm)
-    gm = _clone_folded_constants_at_outputs(gm)
 
     # Delete the constant folder instance which holds GPU memory
     del cf
@@ -101,49 +100,6 @@ def replace_node_with_constant(
     setattr(gm, qualname, constant)
 
 
-def _clone_folded_constants_at_outputs(
-    gm: torch.fx.GraphModule,
-) -> torch.fx.GraphModule:
-    """If a folded buffer is returned from the graph, return a clone instead.
-
-    Eager code after a graph break may mutate outputs in-place. Without a clone,
-    that mutates self._frozen_param* and poisons later calls.
-    """
-    output_node = next(n for n in gm.graph.nodes if n.op == "output")
-    # FX output args are typically a tuple/list of returned values
-    out_args = output_node.args[0]
-    if not isinstance(out_args, (tuple, list)):
-        out_args = (out_args,)
-
-    clone_cache = {}
-    new_outs = []
-    changed = False
-    for out in out_args:
-        if (
-            isinstance(out, torch.fx.Node)
-            and out.op == "get_attr"
-            and str(out.target).startswith("_frozen_param")
-        ):
-            if out not in clone_cache:
-                with gm.graph.inserting_before(output_node):
-                    cloned = gm.graph.call_function(torch.clone, args=(out,))
-                    # keep meta if present
-                    if hasattr(out, "meta"):
-                        cloned.meta.update(out.meta)
-                clone_cache[out] = cloned
-            new_outs.append(clone_cache[out])
-            changed = True
-        else:
-            new_outs.append(out)
-
-    if changed:
-        output_node.args = (tuple(new_outs),)
-        gm.graph.lint()
-        gm.recompile()
-
-    return gm
-
-
 # TODO: Delete this class when the following code is fixed in nightly:
 # https://github.com/pytorch/pytorch/blob/4b881b0da390c1290bb12850ef9daad6f6eb2cb6/torch/_inductor/constant_folding.py#L53-L63
 class _TorchTensorRTConstantFolder(ConstantFolder):  # type: ignore[misc]
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
new file mode 100644
index 0000000000..85562e32d3
--- /dev/null
+++ b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
@@ -0,0 +1,50 @@
+import logging
+from typing import Dict
+
+import torch
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
+    clean_up_graph_after_modifications,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def reset_folded_constructors(
+    gm: torch.fx.GraphModule, settings: CompilationSettings
+) -> torch.fx.GraphModule:
+    """Clone folded constructors that escape through a graph boundary.
+
+    A ``_frozen_param*`` is compiler-owned state, unlike a placeholder supplied
+    by the caller. If a folded constructor becomes an output, eager code or a
+    downstream partition may mutate it. Cloning at that boundary gives each
+    invocation fresh storage while preserving aliases between repeated outputs.
+
+    This pass is intentionally separate from constant folding so it can run
+    again after partitioning, when new TensorRT subgraph outputs are known.
+    """
+    output_node = next(node for node in gm.graph.nodes if node.op == "output")
+    clone_cache: Dict[torch.fx.Node, torch.fx.Node] = {}
+
+    def clone_folded_output(node: torch.fx.Node) -> torch.fx.Node:
+        if node.op != "get_attr" or not str(node.target).startswith("_frozen_param"):
+            return node
+
+        if node not in clone_cache:
+            with gm.graph.inserting_before(output_node):
+                clone = gm.graph.call_function(
+                    torch.ops.aten.clone.default, args=(node,)
+                )
+                clone.meta.update(node.meta)
+            clone_cache[node] = clone
+
+        return clone_cache[node]
+
+    new_output = torch.fx.map_arg(output_node.args[0], clone_folded_output)
+    if not clone_cache:
+        return gm
+
+    output_node.args = (new_output,)
+    gm = clean_up_graph_after_modifications(gm)
+    logger.debug("Reset folded constructors at graph outputs:\n%s", gm.graph)
+    return gm
diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
index b108eb5fa5..85ff9e2974 100644
--- a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
+++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
@@ -11,6 +11,9 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold
+from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import (
+    reset_folded_constructors,
+)
 
 
 @torch._dynamo.disable
@@ -50,7 +53,7 @@ def test_state_is_fresh_across_calls(self):
         torch.testing.assert_close(first, torch.ones_like(first))
         torch.testing.assert_close(second, torch.full_like(second, 3.0))
 
-    def test_constant_fold_clones_frozen_outputs(self):
+    def test_folded_constructor_outputs_are_reset(self):
         """Folded constants returned as graph outputs must be cloned."""
         g = Graph()
         with g.inserting_after():
@@ -68,6 +71,7 @@ def test_constant_fold_clones_frozen_outputs(self):
 
         gm = GraphModule(torch.nn.Module(), g)
         gm = constant_fold(gm, CompilationSettings())
+        gm = reset_folded_constructors(gm, CompilationSettings())
 
         # Outputs should be clones of _frozen_param*, not the attrs themselves.
         output_node = next(n for n in gm.graph.nodes if n.op == "output")
@@ -105,6 +109,7 @@ def test_repeated_folded_outputs_preserve_alias(self):
 
         gm = GraphModule(torch.nn.Module(), g)
         gm = constant_fold(gm, CompilationSettings())
+        gm = reset_folded_constructors(gm, CompilationSettings())
 
         output_node = next(n for n in gm.graph.nodes if n.op == "output")
         out0, out1 = output_node.args[0]
@@ -128,6 +133,44 @@ def test_repeated_folded_outputs_preserve_alias(self):
         for t in frozen:
             self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8)))
 
+    def test_post_partition_frozen_output_is_reset(self):
+        """A frozen value exposed as a subgraph output is cloned by the late pass."""
+        root = torch.nn.Module()
+        root.register_parameter(
+            "_frozen_param0",
+            torch.nn.Parameter(torch.zeros(8), requires_grad=False),
+        )
+        g = Graph()
+        frozen = g.get_attr("_frozen_param0")
+        g.output({"first": frozen, "nested": (frozen,)})
+
+        gm = GraphModule(root, g)
+        gm = reset_folded_constructors(gm, CompilationSettings())
+
+        first = gm()
+        self.assertEqual(first["first"].data_ptr(), first["nested"][0].data_ptr())
+        first["first"].add_(1)
+        self.assertTrue(torch.equal(first["first"], first["nested"][0]))
+        self.assertTrue(torch.equal(gm._frozen_param0, torch.zeros(8)))
+
+        second = gm()
+        self.assertTrue(torch.equal(second["first"], torch.zeros(8)))
+
+    def test_user_owned_output_is_not_reset(self):
+        """Graph inputs remain caller-owned and retain copy-by-reference semantics."""
+        g = Graph()
+        value = g.placeholder("value")
+        g.output(value)
+
+        gm = GraphModule(torch.nn.Module(), g)
+        gm = reset_folded_constructors(gm, CompilationSettings())
+
+        supplied = torch.zeros(8)
+        returned = gm(supplied)
+        self.assertEqual(returned.data_ptr(), supplied.data_ptr())
+        returned.add_(1)
+        self.assertTrue(torch.equal(supplied, torch.ones(8)))
+
 
 if __name__ == "__main__":
     run_tests()

From 0159364491e9eb34518e306188921b9b3a6cef32 Mon Sep 17 00:00:00 2001
From: Michael Williams <mww565@icloud.com>
Date: Thu, 20 Aug 2026 16:47:15 -0700
Subject: [PATCH 4/5] Drop internal frozen-param names from pass comments.

Comments describe folded constructors in plain language instead of Sphinx markup and implementation attribute names.
---
 .../dynamo/lowering/passes/reset_folded_constructors.py  | 8 ++++----
 .../lowering/test_constant_folding_mutable_outputs.py    | 9 ++++-----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
index 85562e32d3..ec83bbe685 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
@@ -15,10 +15,10 @@ def reset_folded_constructors(
 ) -> torch.fx.GraphModule:
     """Clone folded constructors that escape through a graph boundary.
 
-    A ``_frozen_param*`` is compiler-owned state, unlike a placeholder supplied
-    by the caller. If a folded constructor becomes an output, eager code or a
-    downstream partition may mutate it. Cloning at that boundary gives each
-    invocation fresh storage while preserving aliases between repeated outputs.
+    A folded constructor is compiler-owned state, unlike a placeholder supplied
+    by the caller. If it becomes an output, eager code or a downstream partition
+    may mutate it. Cloning at that boundary gives each invocation fresh storage
+    while preserving aliases between repeated outputs.
 
     This pass is intentionally separate from constant folding so it can run
     again after partitioning, when new TensorRT subgraph outputs are known.
diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
index 85ff9e2974..d1e27a6ad2 100644
--- a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
+++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
@@ -1,9 +1,8 @@
 # Regression tests for https://github.com/pytorch/TensorRT/issues/4466
 #
 # Constant folding may replace input-independent factories (e.g. torch.zeros)
-# with persistent _frozen_param* attributes. If those are returned across a
-# graph break and mutated in eager code, later calls must not observe that
-# mutation.
+# with persistent module attributes. If those are returned across a graph break
+# and mutated in eager code, later calls must not observe that mutation.
 
 import torch
 import torch_tensorrt  # noqa: F401  # Registers the "tensorrt" backend
@@ -26,7 +25,7 @@ def mutate_accumulators(values, weight, sample):
 
 def accumulate_with_fresh_state(sample):
     # Input-independent factories; constant folding replaces these with
-    # persistent _frozen_param* attributes.
+    # persistent module attributes.
     values = torch.zeros((8,), dtype=torch.float32, device="cuda")
     weight = torch.zeros((8,), dtype=torch.float32, device="cuda")
     return mutate_accumulators(values, weight, sample)
@@ -73,7 +72,7 @@ def test_folded_constructor_outputs_are_reset(self):
         gm = constant_fold(gm, CompilationSettings())
         gm = reset_folded_constructors(gm, CompilationSettings())
 
-        # Outputs should be clones of _frozen_param*, not the attrs themselves.
+        # Outputs should be clones of the folded constants, not the attrs themselves.
         output_node = next(n for n in gm.graph.nodes if n.op == "output")
         outs = output_node.args[0]
         self.assertEqual(len(outs), 2)

From 20f11e5b88c7f9b69d511d243e76f29d4985a696 Mon Sep 17 00:00:00 2001
From: Michael Williams <mww565@icloud.com>
Date: Fri, 21 Aug 2026 12:26:51 -0700
Subject: [PATCH 5/5] Tag folded constructors at fold time and clone them at
 construction so module state stays persistent while function-local factories
 reset each call.

---
 .../lowering/passes/constant_folding.py       |  6 ++
 .../passes/reset_folded_constructors.py       | 90 ++++++++++++-------
 .../test_constant_folding_mutable_outputs.py  | 78 +++++++++++++++-
 3 files changed, 141 insertions(+), 33 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
index fd4b350825..d7e4ea6bd4 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
@@ -7,6 +7,9 @@
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
     clean_up_graph_after_modifications,
 )
+from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import (
+    FOLDED_CONSTRUCTOR_META,
+)
 
 from packaging import version
 
@@ -93,6 +96,9 @@ def replace_node_with_constant(
         new_input_node = g.create_node("get_attr", qualname, (), {})
         node.replace_all_uses_with(new_input_node)
         new_input_node.meta.update(node.meta)
+        # Distinguishes this from module state that existed before folding: the
+        # value came from an op in the graph body, so eager rebuilds it per call.
+        new_input_node.meta[FOLDED_CONSTRUCTOR_META] = True
         g.erase_node(node)
 
     # Needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
index ec83bbe685..9cd50c58fa 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py
@@ -1,5 +1,4 @@
 import logging
-from typing import Dict
 
 import torch
 from torch_tensorrt.dynamo._settings import CompilationSettings
@@ -9,42 +8,69 @@
 
 logger = logging.getLogger(__name__)
 
+# Set by constant folding on every attribute it creates. Such a value was built
+# by an op in the graph body, so eager semantics allocate it again on every
+# call. Attributes that already existed on the module are never folded, so they
+# never carry this tag and keep their persistent, caller-visible identity.
+FOLDED_CONSTRUCTOR_META = "folded_constructor"
+
+# Marks the copy inserted below, so running this pass again (once per TensorRT
+# submodule after partitioning) does not stack redundant copies.
+_FRESH_COPY_META = "folded_constructor_reset"
+
+
+def _mutates_its_input(user: torch.fx.Node) -> bool:
+    target = user.target
+    if not isinstance(target, torch._ops.OpOverload):
+        return False
+    return bool(target._schema.is_mutable)
+
 
 def reset_folded_constructors(
     gm: torch.fx.GraphModule, settings: CompilationSettings
 ) -> torch.fx.GraphModule:
-    """Clone folded constructors that escape through a graph boundary.
+    """Rebuild folded constructors that must not persist between calls.
+
+    Constant folding hoists ops out of the graph body into module attributes,
+    which turns per-call values into state shared by every invocation. That is
+    only observable when the value escapes as an output or is mutated in place;
+    a folded value that is merely read stays a genuine constant and is left as
+    is, so weights keep converting to TensorRT constants.
 
-    A folded constructor is compiler-owned state, unlike a placeholder supplied
-    by the caller. If it becomes an output, eager code or a downstream partition
-    may mutate it. Cloning at that boundary gives each invocation fresh storage
-    while preserving aliases between repeated outputs.
+    The copy is inserted at the point of construction rather than at the return,
+    so an in-place mutation inside the graph also sees fresh storage. Every user
+    reads the same copy, which preserves aliasing when a value is returned more
+    than once.
 
-    This pass is intentionally separate from constant folding so it can run
-    again after partitioning, when new TensorRT subgraph outputs are known.
+    Values the caller owns, such as placeholders and pre-existing module
+    attributes, are untagged and deliberately untouched: Python passes those by
+    reference and mutations are expected to persist.
     """
-    output_node = next(node for node in gm.graph.nodes if node.op == "output")
-    clone_cache: Dict[torch.fx.Node, torch.fx.Node] = {}
-
-    def clone_folded_output(node: torch.fx.Node) -> torch.fx.Node:
-        if node.op != "get_attr" or not str(node.target).startswith("_frozen_param"):
-            return node
-
-        if node not in clone_cache:
-            with gm.graph.inserting_before(output_node):
-                clone = gm.graph.call_function(
-                    torch.ops.aten.clone.default, args=(node,)
-                )
-                clone.meta.update(node.meta)
-            clone_cache[node] = clone
-
-        return clone_cache[node]
-
-    new_output = torch.fx.map_arg(output_node.args[0], clone_folded_output)
-    if not clone_cache:
-        return gm
-
-    output_node.args = (new_output,)
-    gm = clean_up_graph_after_modifications(gm)
-    logger.debug("Reset folded constructors at graph outputs:\n%s", gm.graph)
+    modified = False
+
+    for node in list(gm.graph.nodes):
+        if node.op != "get_attr" or not node.meta.get(FOLDED_CONSTRUCTOR_META):
+            continue
+
+        users = list(node.users)
+        if not users or any(user.meta.get(_FRESH_COPY_META) for user in users):
+            continue
+        if not any(user.op == "output" or _mutates_its_input(user) for user in users):
+            continue
+
+        with gm.graph.inserting_after(node):
+            fresh = gm.graph.call_function(torch.ops.aten.clone.default, args=(node,))
+        fresh.meta.update(node.meta)
+        fresh.meta.pop(FOLDED_CONSTRUCTOR_META, None)
+        fresh.meta[_FRESH_COPY_META] = True
+
+        for user in users:
+            user.replace_input_with(node, fresh)
+
+        modified = True
+
+    if modified:
+        gm = clean_up_graph_after_modifications(gm)
+        logger.debug("Graph after resetting folded constructors:\n%s", gm.graph)
+
     return gm
diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
index d1e27a6ad2..1ac07394ab 100644
--- a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
+++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py
@@ -11,6 +11,7 @@
 from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold
 from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import (
+    FOLDED_CONSTRUCTOR_META,
     reset_folded_constructors,
 )
 
@@ -133,7 +134,7 @@ def test_repeated_folded_outputs_preserve_alias(self):
             self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8)))
 
     def test_post_partition_frozen_output_is_reset(self):
-        """A frozen value exposed as a subgraph output is cloned by the late pass."""
+        """A folded value exposed as a subgraph output is reset by the late pass."""
         root = torch.nn.Module()
         root.register_parameter(
             "_frozen_param0",
@@ -141,6 +142,7 @@ def test_post_partition_frozen_output_is_reset(self):
         )
         g = Graph()
         frozen = g.get_attr("_frozen_param0")
+        frozen.meta[FOLDED_CONSTRUCTOR_META] = True
         g.output({"first": frozen, "nested": (frozen,)})
 
         gm = GraphModule(root, g)
@@ -170,6 +172,80 @@ def test_user_owned_output_is_not_reset(self):
         returned.add_(1)
         self.assertTrue(torch.equal(supplied, torch.ones(8)))
 
+    def _mutate_and_return_attr(self, root, attr, folded):
+        """Graph for `self.<attr> += 1; return self.<attr>`."""
+        g = Graph()
+        state = g.get_attr(attr)
+        if folded:
+            state.meta[FOLDED_CONSTRUCTOR_META] = True
+        g.output(g.call_function(torch.ops.aten.add_.Tensor, args=(state, 1)))
+        gm = GraphModule(root, g)
+        return reset_folded_constructors(gm, CompilationSettings())
+
+    def test_module_state_keeps_mutation_between_calls(self):
+        """State assigned in __init__ is caller-visible, so mutation must persist."""
+        root = torch.nn.Module()
+        root.register_buffer("weight", torch.zeros(8))
+
+        gm = self._mutate_and_return_attr(root, "weight", folded=False)
+
+        self.assertTrue(torch.equal(gm(), torch.ones(8)))
+        self.assertTrue(torch.equal(gm(), torch.full((8,), 2.0)))
+        self.assertTrue(torch.equal(gm.weight, torch.full((8,), 2.0)))
+
+    def test_function_state_is_rebuilt_between_calls(self):
+        """A constructor folded out of forward must not accumulate across calls."""
+        root = torch.nn.Module()
+        root.register_parameter(
+            "_frozen_param0",
+            torch.nn.Parameter(torch.zeros(8), requires_grad=False),
+        )
+
+        gm = self._mutate_and_return_attr(root, "_frozen_param0", folded=True)
+
+        # The copy precedes the in-place op, so the stored value never changes.
+        self.assertTrue(torch.equal(gm(), torch.ones(8)))
+        self.assertTrue(torch.equal(gm(), torch.ones(8)))
+        self.assertTrue(torch.equal(gm._frozen_param0.detach(), torch.zeros(8)))
+
+    def test_read_only_folded_constant_is_left_alone(self):
+        """A folded constant that only feeds an op stays a constant for TensorRT."""
+        root = torch.nn.Module()
+        root.register_parameter(
+            "_frozen_param0",
+            torch.nn.Parameter(torch.ones(8), requires_grad=False),
+        )
+        g = Graph()
+        weight = g.get_attr("_frozen_param0")
+        weight.meta[FOLDED_CONSTRUCTOR_META] = True
+        value = g.placeholder("value")
+        g.output(g.call_function(torch.ops.aten.add.Tensor, args=(value, weight)))
+
+        gm = GraphModule(root, g)
+        gm = reset_folded_constructors(gm, CompilationSettings())
+
+        clones = [n for n in gm.graph.nodes if n.target is torch.ops.aten.clone.default]
+        self.assertEqual(clones, [])
+
+    def test_repeated_runs_do_not_stack_copies(self):
+        """Re-running per TensorRT submodule must not add a copy each time."""
+        root = torch.nn.Module()
+        root.register_parameter(
+            "_frozen_param0",
+            torch.nn.Parameter(torch.zeros(8), requires_grad=False),
+        )
+        g = Graph()
+        frozen = g.get_attr("_frozen_param0")
+        frozen.meta[FOLDED_CONSTRUCTOR_META] = True
+        g.output(frozen)
+
+        gm = GraphModule(root, g)
+        for _ in range(3):
+            gm = reset_folded_constructors(gm, CompilationSettings())
+
+        clones = [n for n in gm.graph.nodes if n.target is torch.ops.aten.clone.default]
+        self.assertEqual(len(clones), 1)
+
 
 if __name__ == "__main__":
     run_tests()