mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
221350e3a4
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132352 Approved by: https://github.com/ezyang ghstack dependencies: #132335, #132351
68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
# 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()
|