Properly trace into mark_static (#120232)

ezyang · pytorchmergebot · commit 91994684016c · 2024-02-21T13:51:31.000Z
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #120232 Approved by: https://github.com/yanboliang
diff --git a/test/dynamo/test_unspec.py b/test/dynamo/test_unspec.py
@@ -273,7 +273,7 @@ def fn(x):
             comptime.assert_static(x.size(0))
             return x + 1
 
-        opt_fn = torch.compile(fn, dynamic=True)
+        opt_fn = torch.compile(fn, dynamic=True, fullgraph=True)
         opt_fn(torch.randn(12, 23))
 
     def test_shape_graph_break(self):
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -187,6 +187,7 @@
     "torch._C._functorch._add_batch_dim": TorchInGraphFunctionVariable,
     "torch._C._functorch._remove_batch_dim": TorchInGraphFunctionVariable,
     "torch._C._functorch.is_batchedtensor": TorchInGraphFunctionVariable,
+    "torch._dynamo.mark_static": UserFunctionVariable,
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -187,6 +187,7 @@`
`187`	`187`	`"torch._C._functorch._add_batch_dim": TorchInGraphFunctionVariable,`
`188`	`188`	`"torch._C._functorch._remove_batch_dim": TorchInGraphFunctionVariable,`
`189`	`189`	`"torch._C._functorch.is_batchedtensor": TorchInGraphFunctionVariable,`
	`190`	`+ "torch._dynamo.mark_static": UserFunctionVariable,`
`190`	`191`	`}`
`191`	`192`
`192`	`193`