From 8fca2643e70484682e32a5cd156d798dc2611185 Mon Sep 17 00:00:00 2001
From: "zhaozhenlong.zzl" <zhaozhenlong.zzl@alibaba-inc.com>
Date: Wed, 29 Jul 2026 10:55:35 +0800
Subject: [PATCH 1/2] fix(dynamo): clone folded constants returned as graph
 outputs

---
 .../lowering/passes/repair_input_as_output.py | 40 +++++++++++++------
 .../lowering/test_aten_lowering_passes.py     | 27 +++++++++++++
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
index 3773550d99..24102a7373 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
@@ -19,29 +19,43 @@ def repair_input_as_output(
     """
     modified_graph = False
 
-    # Extract graph placeholder Tensors
+    # Extract graph placeholder Tensors. Constant folding also turns
+    # input-independent tensor factories (for example torch.zeros) into
+    # registered `_frozen_param*` attributes. When such a value is returned
+    # across a graph break, eager code may legally mutate it in-place. Return
+    # a clone so the registered constant itself remains immutable.
     placeholders = get_tensor_placeholders(gm)
+    folded_constants = [
+        node
+        for node in gm.graph.nodes
+        if node.op == "get_attr" and str(node.target).startswith("_frozen_param")
+    ]
 
-    for placeholder in placeholders:
-        # If any placeholder has any users which are direct graph outputs
-        if len(placeholder.users) >= 1 and any(
-            user.op == "output" for user in placeholder.users
+    for source in [*placeholders, *folded_constants]:
+        # If any source has any users which are direct graph outputs
+        if len(source.users) >= 1 and any(
+            user.op == "output" for user in source.users
         ):
             modified_graph = True
 
-            # Get direct graph outputs which are direct uses of placeholders
-            direct_outputs = [user for user in placeholder.users if user.op == "output"]
+            # Get graph outputs which directly use the source
+            direct_outputs = [
+                user for user in source.users if user.op == "output"
+            ]
 
-            # Insert clone node for placeholder to ensure placeholder is not a direct output
-            with gm.graph.inserting_after(placeholders[-1]):
-                cloned_placeholder = gm.graph.call_function(
+            # Insert a clone so the source is not returned directly
+            insertion_point = (
+                placeholders[-1] if source.op == "placeholder" else source
+            )
+            with gm.graph.inserting_after(insertion_point):
+                cloned_source = gm.graph.call_function(
                     torch.ops.aten.clone.default,
-                    args=(placeholder,),
+                    args=(source,),
                 )
 
-            # Replace placeholder as output with cloned version
+            # Replace the direct output with the cloned version
             for output in direct_outputs:
-                output.replace_input_with(placeholder, cloned_placeholder)
+                output.replace_input_with(source, cloned_source)
 
     if modified_graph:
         gm = clean_up_graph_after_modifications(gm)
diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
index 51b31c8da1..79fc4a46a8 100644
--- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py
+++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py
@@ -57,6 +57,33 @@ def forward(self, x, y):
         )
         torch._dynamo.reset()
 
+    def test_folded_constant_as_output_is_cloned(self):
+        from torch_tensorrt.dynamo._settings import CompilationSettings
+        from torch_tensorrt.dynamo.lowering.passes.repair_input_as_output import (
+            repair_input_as_output,
+        )
+
+        root = torch.nn.Module()
+        root.register_parameter(
+            "_frozen_param0",
+            torch.nn.Parameter(torch.zeros(8), requires_grad=False),
+        )
+        graph = torch.fx.Graph()
+        folded_constant = graph.get_attr("_frozen_param0")
+        graph.output((folded_constant,))
+        graph_module = torch.fx.GraphModule(root, graph)
+
+        repaired_module = repair_input_as_output(
+            graph_module, CompilationSettings()
+        )
+        first_output = repaired_module()[0]
+        first_output.add_(1)
+        second_output = repaired_module()[0]
+
+        self.assertEqual(first_output, torch.ones(8))
+        self.assertEqual(second_output, torch.zeros(8))
+        self.assertEqual(repaired_module._frozen_param0, torch.zeros(8))
+
 
 class TestLoweringPassMembership(TestCase):
     def insert_at_end(self):

From 673784ef75a42800cfc6c265fe44f00bbc71e99f Mon Sep 17 00:00:00 2001
From: "zhaozhenlong.zzl" <zhaozhenlong.zzl@alibaba-inc.com>
Date: Wed, 29 Jul 2026 15:39:37 +0800
Subject: [PATCH 2/2] fix(dynamo): preserve metadata on cloned outputs

---
 .../dynamo/lowering/passes/repair_input_as_output.py             | 1 +
 1 file changed, 1 insertion(+)

diff --git a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
index 24102a7373..9da585ed5a 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py
@@ -52,6 +52,7 @@ def repair_input_as_output(
                     torch.ops.aten.clone.default,
                     args=(source,),
                 )
+                cloned_source.meta = source.meta.copy()
 
             # Replace the direct output with the cloned version
             for output in direct_outputs: