Graph in symbolic of autograd.Function is not GraphContext

lambAmirstan · November 21, 2022, 10:42am

Hi
Since torch1.13+ torch._C.Graph has been wrapped in GraphContext in the registered symbolic functions.

pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1838


      
          # See Note [Export inplace]
          node_kind = node.kind()
          if node_kind.endswith("_"):
              # Treat relu_ -> relu; add_ -> add etc.
              ns_op_name = node_kind[:-1]
          else:
              ns_op_name = node_kind
          
          
namespace, op_name = ns_op_name.split("::")
          
          
graph_context = jit_utils.GraphContext(
              graph=graph,
              block=block,
              opset=opset_version,
              original_node=node,
              params_dict=_params_dict,
              env=env,
          )
          
          
try:
              # Caffe2-specific: Quantized op symbolics are registered for opset 9 only.

We can get the opset or envs from the context, that’s cool.

But when we export a custom autograd.Function, the graph in the symbolic is still torch._C.Graph. _run_script_method simply run the symbolic without GraphContext wrapping.

github.com

pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1716


      
              set_names(list(graph.outputs()), output_names, "output")
          
          

          
@_beartype.beartype
          def _run_symbolic_method(g, op_name, symbolic_fn, args):
              r"""
              This trampoline function gets invoked for every symbolic method
              call from C++.
              """
              try:
                  return symbolic_fn(g, *args)
              except TypeError as e:
                  # Handle the specific case where we didn't successfully dispatch
                  # to symbolic_fn.  Otherwise, the backtrace will have the clues
                  # you need.
                  e.args = (f"{e.args[0]} (occurred when translating {op_name})",)
                  raise
          
          

          
@_beartype.beartype
          def _add_block(node: _C.Node) -> _C.Block:

We can not reuse the tools in torch.onnx since some of then might branched by the opset or envs, which is not available in torch._C.Graph.

https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/symbolic_opset9.py#L5335

Is it expected? Or should PyTorch create a GraphContext in _run_script_method?