pytorch
3d231667 - [ONNX] Create GraphContext and load `g.op` method to the class (#84728)

Commit
2 years ago
[ONNX] Create GraphContext and load `g.op` method to the class (#84728) This PR create the `GraphContext` class and relays all graph methods to _C.Graph as well as implements the `g.op` method. The GraphContext object is passed into the symbolic functions in place of _C.Graph for compatibility with existing symbolic functions. This way (1) we can type annotate all `g` args because the method is defined and (2) we can use additional context information in symbolic functions. (3) no more monkey patching on `_C.Graph` Also - Fix return type of `_jit_pass_fixup_onnx_controlflow_node` - Create `torchscript.py` to house torch.Graph related functions - Change `GraphContext.op` to create nodes in the Block instead of the Graph - Create `add_op_with_blocks` to handle scenarios where we need to directly manipulate sub-blocks. Update loop and if symbolic functions to use this function. ## Discussion Should we put all the context inside `SymbolicContext` and make it an attribute in the `GraphContext` class? This way we only define two attributes `GraphContext.graph` and `GraphContext.context`. Currently all context attributes are directly defined in the class. ### Decision Keep GraphContext flatand note that it will change in the future. Pull Request resolved: https://github.com/pytorch/pytorch/pull/84728 Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
Author
Committer
Parents
Loading