class_docstr: A backend compiles (sub-)graphs in torch.fx.GraphModule to onnxruntime.InferenceSession calls.
  
  The compiler entry point is OrtBackend.compile, which
      1. partitions the original graph into supported sub-graphs (type: torch.fx.GraphModule) and unsupported
         sub-graphs.
      2. For each supported sub-graph, it replaces its _wrapped_call function with _ort_accelerated_call.
      3. Inside _ort_accelerated_call, it creates onnxruntime.InferenceSession and calls it to execute the sub-graph.
class_docstr_tok: ["A","backend","compiles","(","sub-",")","graphs","in","torch.fx.GraphModule","to","onnxruntime.InferenceSession","calls",".","The","compiler","entry","point","is","OrtBackend.compile",",","which","1.","partitions","the","original","graph","into","supported","sub-graphs","(","type",":","torch.fx.GraphModule",")","and","unsupported","sub-graphs",".","2",".","For","each","supported","sub-graph",",","it","replaces","its","_wrapped_call","function","with","_ort_accelerated_call",".","3",".","Inside","_ort_accelerated_call",",","it","creates","onnxruntime.InferenceSession","and","calls","it","to","execute","the","sub-graph","."]
code: def compile(
      self, graph_module: torch.fx.GraphModule, args
  ) -> torch.fx.GraphModule:
      # Deferred import since CapabilityBasedPartitioner is not decorated with
      # @compatibility; importing it at the module level will result in the test
      # failing: pytest test/test_fx.py -k test_public_api_surface
      # because this module is imported into torch.onnx.
      from torch.fx.passes.infra.partitioner import (
          CapabilityBasedPartitioner,
      )
  
      # FX graph based partitioning based on ONNX supported ops.
      # Given a graph module
      #  GraphModule0
      #   node_0
      #   node_1
      #   node_2
      #   node_3
      #   node_4
      # If only node_2 is not supported by ONNX, this graph module will be partitioned into
      #  GraphModule0
      #   GraphModule1
      #    node_0
      #    node_1
      #   node_2
      #   GraphModule2
      #    node_3
      #    node_4
      # by calling CapabilityBasedPartitioner.partition_and_fuse.
      # Then, GraphModule1's and GraphModule2's forward method (GraphModule._wrapped_call)
      # will be replaced by OrtBackend._ort_accelerated_call to delegate computation to ORT.
      if graph_module in self._partitioner_cache:
          partitioned_prim_graph_module = self._partitioner_cache[
              graph_module
          ]
      else:
          prim_graph_module = graph_module
          partitioner = CapabilityBasedPartitioner(
              prim_graph_module,
              self._supported_ops,
              allows_single_node_partition=True,
          )
          partitioned_prim_graph_module = (
              partitioner.partition_and_fuse()
          )
          self._partitioner_cache[
              graph_module
          ] = partitioned_prim_graph_module
  
          # Overriding fused_module's __call__() function with ort_acclerated_call()
          # This loop goes through all graph partitions (each of them is an ONNX-representable graph)
          # and override their _wrapped_call function with _ort_accelerated_call.
          # Inside _ort_accelerated_call, the partition's graph is exported into ONNX and executed by ORT.
          for node in partitioned_prim_graph_module.graph.nodes:
              # TODO(wschin): use a better way to identify fused submodule
              # See https://github.com/pytorch/pytorch/issues/106872.
              if node.op == "call_module" and "fused_" in node.name:
                  fused_module = getattr(
                      partitioned_prim_graph_module, node.name
                  )
                  # self.ort_acclerated_call is responsible for exporting graph to ONNX,
                  # creating ORT session, and running ORT session.
                  fused_module._wrapped_call = self._ort_acclerated_call
  
      return partitioned_prim_graph_module
  
imports: import dataclasses
  import importlib
  import logging
  from typing import Any, Dict, Final, List, Mapping, Optional, Sequence, Set, Tuple, Union
  from typing_extensions import TypeAlias
  import torch
  import torch._C
  import torch._ops
  import torch._prims.executor
  import torch.fx
  from torch._subclasses.fake_tensor import FakeTensor
  from torch.fx._compatibility import compatibility
  from torch.fx.passes.fake_tensor_prop import FakeTensorProp
  from torch.fx.passes.operator_support import OperatorSupport
  from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
  from torch.utils import _pytree
label_desc: Use image node_id 4 for calling the OrtBackend obj's underlying member method code with example usage: obj.compile(graph_module, args) and returns: partitioned_prim_graph_module