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","."]
docstr: This function replaces GraphModule._wrapped_call in compiled model.
  
  The _wrapped_call is the underlying implementation of forward method. Replacing
  it means we delegate the computation to _ort_acclerated_call and therefore
  onnxruntime.InferenceSession.
code: def _ort_acclerated_call(
      self, graph_module: torch.fx.GraphModule, *args, **kwargs
  ):
      """This function replaces GraphModule._wrapped_call in compiled model.
  
      The _wrapped_call is the underlying implementation of forward method. Replacing
      it means we delegate the computation to _ort_acclerated_call and therefore
      onnxruntime.InferenceSession.
      """
      cached_execution_info_per_session = self._all_ort_execution_info.search_reusable_session_execution_info(
          graph_module, *args
      )
      if cached_execution_info_per_session:
          onnx_session = cached_execution_info_per_session.session
          input_names = cached_execution_info_per_session.input_names
          output_names = cached_execution_info_per_session.output_names
          input_value_infos = (
              cached_execution_info_per_session.input_value_infos
          )
          output_value_infos = (
              cached_execution_info_per_session.output_value_infos
          )
          input_devices = (
              cached_execution_info_per_session.input_devices
          )
          output_devices = (
              cached_execution_info_per_session.output_devices
          )
          prim_outputs = (
              cached_execution_info_per_session.example_outputs
          )
      else:
          # It's first time seeing such as graph. Let's make a new session
          # (type: onnxruntime.InferenceSession) for it.
  
          graph_module = torch.onnx._internal.fx.passes.MovePlaceholderToFront(
              self._resolved_onnx_exporter_options.diagnostic_context,
              graph_module,
          ).run()
          # Generate reference outputs. They are used to indicate output
          # tensors' types and devices when calling ORT.
          #
          # WARNING: The downstream code should not change prim_outputs and
          # this backend should always produces output with schema identical to prim_outputs'.
  
          if self._resolved_onnx_exporter_options.dynamic_shapes:
              # No pre-allocation when dynamic shape is enabled.
              self.preallocate_output = False
              extracted_outputs = _extract_graph_module_outputs(
                  graph_module
              )
  
              def maybe_map_to_meta_val(value):
                  if hasattr(value, "meta") and "val" in value.meta:
                      # Select outputs with "val" information. Without "val",
                      # it's not possible access output_arg.meta["val"].device.
                      return value.meta["val"]
                  else:
                      return value
  
              prim_outputs = _pytree.tree_map(
                  maybe_map_to_meta_val, extracted_outputs
              )
          else:
              try:
                  prim_outputs = FakeTensorProp(graph_module).propagate(
                      *args, **kwargs
                  )
              except Exception:
                  logger.warning(
                      "FakeTensorProb failed for %s", graph_module
                  )
                  # When FakeTensorProp fails, it is not possible to preallocate output buffers
                  # because the output shapes are not inferred.
                  self.preallocate_output = False
  
                  # rethrow FakeTensorProb failure because it is not yet currently handled.
                  raise
  
          # Create the object to iterate through the nodes in graph one-by-one
          # and calls the corresponding ONNX exporter for each node.
          fx_interpreter = fx_onnx_interpreter.FxOnnxInterpreter(
              diagnostic_context=self._resolved_onnx_exporter_options.diagnostic_context
          )
          # Cast FX variables if they will result schema-mismatch when searching
          # for ONNX operator. E.g., add(double_tensor, int_tensor) is fine in PyTorch,
          # but ONNX expects add(double_tensor, double_tensor).
          graph_module = torch.onnx._internal.fx.passes.InsertTypePromotion(
              self._resolved_onnx_exporter_options.diagnostic_context,
              graph_module,
          ).run()
          # Start the per-node exporting process. It's conceptually a for loop
          # scanning through the nodes in the graph.
          exported = fx_interpreter.run(
              fx_graph_module=graph_module,
              onnxfunction_dispatcher=self._resolved_onnx_exporter_options.onnxfunction_dispatcher,
              op_level_debug=self._resolved_onnx_exporter_options.op_level_debug,
          )
          # Convert the exported result to ONNX ModelProto.
          onnx_model = exported.to_model_proto(
              opset_version=self._resolved_onnx_exporter_options.onnx_registry.opset_version,
          )
  
          # Initialize a ORT session to execute this ONNX model.
          # Note that TorchDynamo assumes all inputs/outputs are on the
          # same device, but it's subject to change (very likely with
          # dynamic shape support), so we add execution providers
          # based on the logic in _select_eps: (explicitly preferred EPs,
          # EPs inferred from inputs or graph, and the fallback default EP)/
          #
          # TODO(wschin): enable external allocators.
          # See https://github.com/pytorch/pytorch/issues/106867
          onnx_session = onnxruntime.InferenceSession(
              path_or_bytes=onnx_model.SerializeToString(),
              sess_options=self._options.ort_session_options,
              providers=self._select_eps(graph_module, *args),
          )
  
          # Cache ORT session. It's reused for the same "graph_module".
          # Generate ONNX model and extract its input and output names.
          input_names = tuple(
              input.name for input in onnx_model.graph.input
          )
          output_names = tuple(
              output.name for output in onnx_model.graph.output
          )
          input_devices = _get_onnx_devices(args)
          # Cache devices for inputs and outputs. They are used to invoke
          # ORT session. Output devices indicate where (e.g., GPU or CPU)
          # to store outputs
          if isinstance(prim_outputs, tuple):
              output_devices = _get_onnx_devices(prim_outputs)
          else:
              output_devices = _get_onnx_devices((prim_outputs,))
  
          input_value_infos = tuple(
              input for input in onnx_model.graph.input
          )
          output_value_infos = tuple(
              output for output in onnx_model.graph.output
          )
  
          execution_info_per_session = OrtExecutionInfoPerSession(
              session=onnx_session,
              input_names=input_names,
              input_value_infos=input_value_infos,
              output_names=output_names,
              output_value_infos=output_value_infos,
              input_devices=input_devices,
              output_devices=output_devices,
              example_outputs=prim_outputs,
          )
  
          self._all_ort_execution_info.cache_session_execution_info(
              graph_module, execution_info_per_session
          )
  
      self.execution_count += 1
  
      # ORT always returns a tuple of outputs. If the original output is a tensor,
      # ORT output's first element must be extracted and returned. Otherwise, type
      # mismatch may happen in downstream computation.
      is_single_tensor_output = isinstance(prim_outputs, torch.Tensor)
      normalized_prim_outputs = (
          (prim_outputs,) if is_single_tensor_output else prim_outputs
      )
      assert isinstance(normalized_prim_outputs, tuple)
      assert all(
          isinstance(elem, (torch.Tensor, torch.SymInt, int))
          for elem in normalized_prim_outputs
      )
  
      _nvtx_range_push("run_onnx_session_with_ortvaluevector")
      onnx_outputs = self.run(
          onnx_session,
          input_names,
          args,
          input_devices,
          output_names,
          normalized_prim_outputs,
          output_devices,
          self._options.preallocate_output,
          input_value_infos,
          normalized_prim_outputs,
      )
      _nvtx_range_pop()
  
      if self._assert_allclose_to_baseline:
          # Compute baseline.
          baseline_outputs = torch._prims.executor.execute(
              graph_module, *args, executor="aten"
          )
          normalized_baseline_ouptuts = (
              (baseline_outputs,)
              if is_single_tensor_output
              else baseline_outputs
          )
          # Ensure every output tensor is close to the corresponding baseline.
          for onnx_output, baseline_output in zip(
              onnx_outputs, normalized_baseline_ouptuts
          ):
              torch.testing.assert_close(onnx_output, baseline_output)
      return (
          onnx_outputs[0] if is_single_tensor_output else onnx_outputs
      )
  
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 3 for calling the OrtBackend obj's underlying member method code with example usage: obj._ort_acclerated_call(graph_module) and returns: unknown, value, value