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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
# Owner(s): ["oncall: export"]
import torch
from torch._dynamo.test_case import TestCase
from torch._export.tools import report_exportability
from torch.testing._internal.common_utils import run_tests
torch.library.define(
"testlib::op_missing_meta",
"(Tensor(a!) x, Tensor(b!) z) -> Tensor",
tags=torch.Tag.pt2_compliant_tag,
)
@torch.library.impl("testlib::op_missing_meta", "cpu")
@torch._dynamo.disable
def op_missing_meta(x, z):
x.add_(5)
z.add_(5)
return x + z
class TestExportTools(TestCase):
def test_report_exportability_basic(self):
class Module(torch.nn.Module):
def forward(self, x, y):
return x[0] + y
f = Module()
inp = ([torch.ones(1, 3)], torch.ones(1, 3))
report = report_exportability(f, inp)
self.assertTrue(len(report) == 1)
self.assertTrue(report[""] is None)
def test_report_exportability_with_issues(self):
class Unsupported(torch.nn.Module):
def forward(self, x):
return torch.ops.testlib.op_missing_meta(x, x.cos())
class Supported(torch.nn.Module):
def forward(self, x):
return x.sin()
class Module(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.unsupported = Unsupported()
self.supported = Supported()
def forward(self, x):
y = torch.nonzero(x)
return self.unsupported(y) + self.supported(y)
f = Module()
inp = (torch.ones(4, 4),)
report = report_exportability(f, inp, strict=False, pre_dispatch=True)
self.assertTrue(report[""] is not None)
self.assertTrue(report["unsupported"] is not None)
self.assertTrue(report["supported"] is None)
if __name__ == "__main__":
run_tests()
|
