From 65ff880f29bed4e096d6804d57a33323d8e9f59c Mon Sep 17 00:00:00 2001
From: john-rocky <rockyshikoku@gmail.com>
Date: Wed, 2 Sep 2026 04:55:43 +0900
Subject: [PATCH] Skip impure ops in constant_prop_pass

constant_prop_pass folds any call_function node whose arguments are all
constants. Ops that draw from the RNG, such as aten.rand, take only sizes
as arguments, so they qualified and were replaced by a single frozen draw:
a model returning x + torch.rand(4) produced the same output on every call
once the pass had run.

Skip nodes that torch.fx.Node.is_impure() reports as impure. That covers
ops tagged nondeterministic_seeded, mutable schemas and side-effectful
functions, and is the same test eliminate_dead_code uses to decide what it
must keep.
---
 exir/passes/constant_prop_pass.py |  4 ++++
 exir/tests/test_passes.py         | 25 +++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py
index 11640d875c0..ea8ee1ad3a9 100644
--- a/exir/passes/constant_prop_pass.py
+++ b/exir/passes/constant_prop_pass.py
@@ -146,6 +146,10 @@ def get_propagated_const_tensor_dict(
             node.op != "call_function"
             or node.target is memory.alloc
             or node.target in all_skip_targets
+            # Ops with side effects (RNG draws, mutation) have to run at
+            # runtime. `aten.rand` has no tensor inputs, so without this check
+            # it would be folded into a single frozen draw.
+            or node.is_impure()
         ):
             continue
 
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 0b586bd44cd..54211a490a3 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -2475,6 +2475,31 @@ def forward(self, x):
         # 1 constant: a (= self.w @ self.cst)
         self.assertEqual(1, len(pass_result.constants))
 
+    def test_constant_prop_pass_skips_nondeterministic_ops(self) -> None:
+        """
+        Ops that draw from the RNG take no tensor inputs, so they look constant
+        to the pass. They have to stay in the graph: folding one would freeze a
+        single random draw into the program.
+        """
+
+        class RandomAdd(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x + torch.rand(4)
+
+        x = torch.zeros(4)
+        edge = to_edge(export(RandomAdd(), (x,), strict=True))
+        new_ep = constant_prop_pass(edge.exported_program())
+
+        rand_nodes = [
+            node
+            for node in new_ep.graph.nodes
+            if node.target == exir_ops.edge.aten.rand.default
+        ]
+        self.assertEqual(len(rand_nodes), 1)
+        self.assertEqual(len(new_ep.constants), 0)
+        module = new_ep.module()
+        self.assertFalse(torch.equal(module(x), module(x)))
+
     def test_constant_prop_pass_zero_stride_tensors(self) -> None:
         """
         Test that constant propagation correctly handles tensors with zero strides