1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
# Owner(s): ["module: dynamo"]
import torch
import torch._dynamo.test_case
def fn_creator():
var1 = 1
def fn(x):
x = x + 1
var2 = 1
torch._dynamo.graph_break()
x = x + var1
def inner_fn():
return var2
return x
return fn
class ResumeFunctionTests(torch._dynamo.test_case.TestCase):
def test_freevars(self):
fn = fn_creator()
opt_fn = torch.compile(fn, backend="eager")
opt_fn(torch.randn(10))
codes = [v for k, v in list(globals().items()) if k.startswith("__resume_at")]
self.assertEqual(len(codes), 1)
# co_freevars of resume functions, are sorted concatenation of the original function's co_freevars and co_cellvars
self.assertEqual(codes[0].co_freevars, ("var1", "var2"))
if __name__ == "__main__":
from torch._dynamo.test_case import run_tests
run_tests()
|
