mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
bf7e290854
This PR is part of a series attempting to re-submit https://github.com/pytorch/pytorch/pull/134592 as smaller PRs. In jit tests: - Add and use a common raise_on_run_directly method for when a user runs a test file directly which should not be run this way. Print the file which the user should have run. - Raise a RuntimeError on tests which have been disabled (not run) Pull Request resolved: https://github.com/pytorch/pytorch/pull/154725 Approved by: https://github.com/clee2000
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
# Owner(s): ["oncall: jit"]
|
|
|
|
import torch
|
|
from torch.testing import FileCheck
|
|
from torch.testing._internal.common_utils import raise_on_run_directly
|
|
from torch.testing._internal.jit_utils import JitTestCase
|
|
|
|
|
|
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))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise_on_run_directly("test/test_jit.py")
|