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
|
# Owner(s): ["module: dynamo"]
from unittest import mock
import torch
import torch._dynamo
import torch._dynamo.test_case
from torch._inductor.utils import pass_execution_and_save
class FxPassesPreGradTests(torch._dynamo.test_case.TestCase):
@mock.patch("torch._inductor.utils.ShapeProp.propagate")
def test_pass_execution_and_save(self, mock_shape_prop):
class TestModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.param = torch.nn.Parameter(torch.ones(4, 4))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.param + x
def fx_pass(graph: torch.fx.GraphModule) -> None:
return
sample_input = torch.randn(4, 4)
m = TestModule()
m(sample_input)
exported_program = torch.export.export(m, (sample_input,))
gm = exported_program.graph_module
pass_execution_and_save(fx_pass, gm, sample_input, "Apply testing pass")
mock_shape_prop.assert_called_once()
if __name__ == "__main__":
from torch._dynamo.test_case import run_tests
run_tests()
|
