python.object-model#
model_attr_mutation#
Original source code:
# mypy: allow-untyped-defs
import torch
class ModelAttrMutation(torch.nn.Module):
"""
Attribute mutation raises a warning. Covered in the test_export.py test_detect_leak_strict test.
"""
def __init__(self) -> None:
super().__init__()
self.attr_list = [torch.randn(3, 2), torch.randn(3, 2)]
def recreate_list(self):
return [torch.zeros(3, 2), torch.zeros(3, 2)]
def forward(self, x):
self.attr_list = self.recreate_list()
return x.sum() + self.attr_list[0].sum()
example_args = (torch.randn(3, 2),)
tags = {"python.object-model"}
model = ModelAttrMutation()
torch.export.export(model, example_args)
Result:
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, x: "f32[3, 2]"):
zeros: "f32[3, 2]" = torch.ops.aten.zeros.default([3, 2], device = device(type='cpu'), pin_memory = False)
zeros_1: "f32[3, 2]" = torch.ops.aten.zeros.default([3, 2], device = device(type='cpu'), pin_memory = False); zeros_1 = None
sum_1: "f32[]" = torch.ops.aten.sum.default(x); x = None
sum_2: "f32[]" = torch.ops.aten.sum.default(zeros); zeros = None
add: "f32[]" = torch.ops.aten.add.Tensor(sum_1, sum_2); sum_1 = sum_2 = None
return (add,)
Graph signature:
# inputs
x: USER_INPUT
# outputs
add: USER_OUTPUT
Range constraints: {}
optional_input#
Original source code:
# mypy: allow-untyped-defs
import torch
from torch._export.db.case import SupportLevel
class OptionalInput(torch.nn.Module):
"""
Tracing through optional input is not supported yet
"""
def forward(self, x, y=torch.randn(2, 3)):
if y is not None:
return x + y
return x
example_args = (torch.randn(2, 3),)
tags = {"python.object-model"}
support_level = SupportLevel.SUPPORTED
model = OptionalInput()
torch.export.export(model, example_args)
Result:
Unsupported: Tracing through optional input is not supported yet
