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): ["oncall: jit"]
import torch
from torch.testing import FileCheck
from torch.testing._internal.jit_utils import JitTestCase
if __name__ == '__main__':
raise RuntimeError("This test file is not meant to be run directly, use:\n\n"
"\tpython test/test_jit.py TESTNAME\n\n"
"instead.")
class TestOpDecompositions(JitTestCase):
def test_op_decomposition(self):
def foo(x):
return torch.var(x, unbiased=True)
# TODO: more robust testing
foo_s = torch.jit.script(foo)
FileCheck().check("aten::var").run(foo_s.graph)
torch._C._jit_pass_run_decompositions(foo_s.graph)
inp = torch.rand([10, 10])
self.assertEqual(foo(inp), foo_s(inp))
FileCheck().check_not("aten::var").run(foo_s.graph)
def test_registered_decomposition(self):
@torch.jit.script
def foo(x):
return torch.square(x)
@torch.jit.script
def square_decomp(x):
return torch.pow(x, 2)
torch.jit._register_decomposition(torch.ops.aten.square.default, square_decomp.graph)
torch._C._jit_pass_run_decompositions(foo.graph)
FileCheck().check_not("aten::square").check("aten::pow").run(foo.graph)
x = torch.rand([4])
self.assertEqual(foo(x), torch.square(x))
|
