Graph in symbolic of autograd.Function is not GraphContext

lambAmirstan: <p>Hi<br>
Since torch1.13+ <code>torch._C.Graph</code> has been wrapped in GraphContext in the registered symbolic functions.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1838">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1838" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1838" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1838</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="1828" style="counter-reset: li-counter 1827 ;">
          <li># See Note [Export inplace]</li>
          <li>node_kind = node.kind()</li>
          <li>if node_kind.endswith("_"):</li>
          <li>    # Treat relu_ -&gt; relu; add_ -&gt; add etc.</li>
          <li>    ns_op_name = node_kind[:-1]</li>
          <li>else:</li>
          <li>    ns_op_name = node_kind</li>
          <li>
          </li>
<li>namespace, op_name = ns_op_name.split("::")</li>
          <li>
          </li>
<li class="selected">graph_context = jit_utils.GraphContext(</li>
          <li>    graph=graph,</li>
          <li>    block=block,</li>
          <li>    opset=opset_version,</li>
          <li>    original_node=node,</li>
          <li>    params_dict=_params_dict,</li>
          <li>    env=env,</li>
          <li>)</li>
          <li>
          </li>
<li>try:</li>
          <li>    # Caffe2-specific: Quantized op symbolics are registered for opset 9 only.</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>We can get the opset or envs from the context, that’s cool.</p>
<p>But when we export a custom autograd.Function, the graph in the symbolic is still <code>torch._C.Graph</code>. <code>_run_script_method</code> simply run the symbolic without GraphContext wrapping.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1716">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1716" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1716" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/utils.py#L1716</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="1706" style="counter-reset: li-counter 1705 ;">
          <li>    set_names(list(graph.outputs()), output_names, "output")</li>
          <li>
          </li>
<li>
          </li>
<li>@_beartype.beartype</li>
          <li>def _run_symbolic_method(g, op_name, symbolic_fn, args):</li>
          <li>    r"""</li>
          <li>    This trampoline function gets invoked for every symbolic method</li>
          <li>    call from C++.</li>
          <li>    """</li>
          <li>    try:</li>
          <li class="selected">        return symbolic_fn(g, *args)</li>
          <li>    except TypeError as e:</li>
          <li>        # Handle the specific case where we didn't successfully dispatch</li>
          <li>        # to symbolic_fn.  Otherwise, the backtrace will have the clues</li>
          <li>        # you need.</li>
          <li>        e.args = (f"{e.args[0]} (occurred when translating {op_name})",)</li>
          <li>        raise</li>
          <li>
          </li>
<li>
          </li>
<li>@_beartype.beartype</li>
          <li>def _add_block(node: _C.Node) -&gt; _C.Block:</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>We can not reuse the tools in <code>torch.onnx</code> since some of then might branched by the opset or envs, which is not available in <code>torch._C.Graph</code>.</p>
<p><a href="https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/symbolic_opset9.py#L5335" class="onebox" target="_blank" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/7c98e70d44abc7a1aead68b6ea6c8adc8c554db5/torch/onnx/symbolic_opset9.py#L5335</a></p>
<p>Is it expected? Or should PyTorch create a GraphContext in <code>_run_script_method</code>?</p>