python.context-manager

null_context_manager

Note

Tags: python.context-manager

Support Level: SUPPORTED

Original source code:

import contextlib

import torch



def null_context_manager(x):
    """
    Null context manager in Python will be traced out.
    """
    ctx = contextlib.nullcontext()
    with ctx:
        return x.sin() + x.cos()

Result:

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[3, 2]):
            #
            sin: f32[3, 2] = torch.ops.aten.sin.default(arg0_1)
            cos: f32[3, 2] = torch.ops.aten.cos.default(arg0_1);  arg0_1 = None
            add: f32[3, 2] = torch.ops.aten.add.Tensor(sin, cos);  sin = cos = None
            return (add,)

Graph Signature: ExportGraphSignature(parameters=[], buffers=[], user_inputs=['arg0_1'], user_outputs=['add'], inputs_to_parameters={}, inputs_to_buffers={}, buffers_to_mutate={}, backward_signature=None, assertion_dep_token=None)
Symbol to range: {}