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
|
# Owner(s): ["oncall: distributed"]
import torch
import torch.nn as nn
from torch.distributed.tensor.debug._op_coverage import get_inductor_decomp_graphs
from torch.testing._internal.common_utils import run_tests, TestCase
class SimpleMLP(nn.Module):
def __init__(self) -> None:
super().__init__()
self.net1 = nn.Linear(50, 32)
self.relu = nn.ReLU()
self.net2 = nn.Linear(32, 8)
def forward(self, x):
return torch.sigmoid(self.net2(self.relu(self.net1(x))))
class TestOpCoverage(TestCase):
def test_trace_with_inductor_decomp(self):
model = SimpleMLP()
args = (torch.randn(8, 50),)
kwargs = {}
graphs = get_inductor_decomp_graphs(model, args, kwargs)
assert len(graphs) == 2, "Expect fwd + bwd graphs"
self.assertIsInstance(graphs[0], torch.fx.GraphModule)
self.assertIsInstance(graphs[1], torch.fx.GraphModule)
if __name__ == "__main__":
run_tests()
|
