torch.mutation¶
user_input_mutation¶
Original source code:
# mypy: allow-untyped-defs
import torch
class UserInputMutation(torch.nn.Module):
"""
Directly mutate user input in forward
"""
def forward(self, x):
x.mul_(2)
return x.cos()
Result:
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, x: "f32[3, 2]"):
mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, 2); x = None
cos: "f32[3, 2]" = torch.ops.aten.cos.default(mul)
return (mul, cos)
Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_INPUT_MUTATION: 6>, arg=TensorArgument(name='mul'), target='x'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='cos'), target=None)])
Range constraints: {}