# mypy: allow-untyped-defsimportcontextlibimporttorchclassNullContextManager(torch.nn.Module):"""    Null context manager in Python will be traced out.    """defforward(self,x):"""        Null context manager in Python will be traced out.        """ctx=contextlib.nullcontext()withctx:returnx.sin()+x.cos()example_args=(torch.randn(3,2),)tags={"python.context-manager"}model=NullContextManager()torch.export.export(model,example_args)

ExportedProgram:classGraphModule(torch.nn.Module):defforward(self,x:"f32[3, 2]"):sin:"f32[3, 2]"=torch.ops.aten.sin.default(x)cos:"f32[3, 2]"=torch.ops.aten.cos.default(x);x=Noneadd:"f32[3, 2]"=torch.ops.aten.add.Tensor(sin,cos);sin=cos=Nonereturn(add,)Graphsignature:# inputsx:USER_INPUT# outputsadd:USER_OUTPUTRangeconstraints:{}