Skip to content

Commit 3d8b4e8

Browse files
committed
switch over decorator usage of dynamo to torch.compile in tests and add inductor test
ghstack-source-id: 1a2e92b Pull Request resolved: #298
1 parent cef6cac commit 3d8b4e8

File tree

1 file changed

+20
-9
lines changed

1 file changed

+20
-9
lines changed

multipy/runtime/test_compat.py

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import unittest
88

99
import torch
10+
import torch._dynamo
1011

1112

1213
class TestCompat(unittest.TestCase):
@@ -22,31 +23,41 @@ def test_pytorch3d(self):
2223
def test_hf_tokenizers(self):
2324
import tokenizers # noqa: F401
2425

25-
@unittest.skip("torch.Library is not supported")
2626
def test_torchdynamo_eager(self):
27-
import torch._dynamo as torchdynamo
2827

29-
@torchdynamo.optimize("eager")
28+
torch._dynamo.reset()
29+
3030
def fn(x, y):
3131
a = torch.cos(x)
3232
b = torch.sin(y)
3333
return a + b
3434

35-
fn(torch.randn(10), torch.randn(10))
35+
c_fn = torch.compile(fn, backend="eager")
36+
c_fn(torch.randn(10), torch.randn(10))
3637

37-
@unittest.skip("torch.Library is not supported")
3838
def test_torchdynamo_ofi(self):
39-
import torch._dynamo as torchdynamo
4039

41-
torchdynamo.reset()
40+
torch._dynamo.reset()
41+
42+
def fn(x, y):
43+
a = torch.cos(x)
44+
b = torch.sin(y)
45+
return a + b
46+
47+
c_fn = torch.compile(fn, backend="ofi")
48+
c_fn(torch.randn(10), torch.randn(10))
49+
50+
def test_torchdynamo_inductor(self):
51+
52+
torch._dynamo.reset()
4253

43-
@torchdynamo.optimize("ofi")
4454
def fn(x, y):
4555
a = torch.cos(x)
4656
b = torch.sin(y)
4757
return a + b
4858

49-
fn(torch.randn(10), torch.randn(10))
59+
c_fn = torch.compile(fn)
60+
c_fn(torch.randn(10), torch.randn(10))
5061

5162

5263
if __name__ == "__main__":

0 commit comments

Comments
 (0)