code: def ExecuteGraph(  # pylint: disable=invalid-name
      pcoll: beam.pvalue.PCollection,
      remote_op_name: str,
      remote_op_name_to_graph_name: Mapping[str, str],
      graph_name_to_specs: Mapping[
          str, List[execution_spec.ExecutionSpec]
      ],
      graph_to_remote_op_input_name_mapping: Mapping[
          str, Mapping[str, Mapping[str, str]]
      ],
  ) -> beam.pvalue.PCollection:
      """A PTransform that executes a graph.
  
      Each graph has a list of ExecutionSpecs, in which the order of the list
      represents the order of execution. An ExecutionSpec can either represent
      a subgraph layer or a remote op in a remote op layer. When executing a
      subgraph layer, we can load and execute the subgraph with a beam ParDo.
      When executing a remote op (which represents another graph), we need to
      load the remote graph inputs, call ExecuteGraph to recursively execute that
      graph, and extract the remote graph output. When executing a remote op, we
      call the current graph "parent" and the remote graph "child".
  
      Here, each Beam element is a dictionary from remote op names to a dictionary
      from tensor names to values, or {remote op name: {tensor name: value}}.
      Note that at any time, PColl only stores input tensor values and computed
      tensor values. The input PColl should have the input tensor names and values
      for the graph ready. As we execute the partitioned subgraphs, we add the
      intermediate output names and values to PColl.
  
      Args:
        pcoll: A PCollection of inputs to the graph. Each element is a dictionary
          from remote op names to a dictionary from tensor names to values. Here,
          element[remote_op_name] contains graph inputs.
        remote_op_name: The remote op name of the current graph.
        remote_op_name_to_graph_name: A mapping from remote op names to graph names.
        graph_name_to_specs: A mapping from graph names to a list of ExecutionSpecs,
          where the order of the list represents the order of execution.
        graph_to_remote_op_input_name_mapping: A mapping from graph names to remote
          op names to remote graph placeholder names to parent graph input names. We
          don't have this information since it was stored in PyFunc's function.
          {graph name: {remote op name: {placeholder name: input name}}}.
  
      Returns:
        A PCollection of results of this graph. Each element is a dictionary from
        remote op names to a dictionary from tensor names to values. Here,
        element[remote_op_name] stores graph inputs, intermediate results, and
        graph outputs.
      """
      specs = graph_name_to_specs[
          remote_op_name_to_graph_name[remote_op_name]
      ]
  
      for spec in specs:
          # Construct Beam subgraph for a subgraph layer.
          if not spec.is_remote_op:
              step_name = "SubgraphLayerDoFn[Graph_%s][Outputs_%s]" % (
                  remote_op_name,
                  "_".join(spec.output_names),
              )
              pcoll = pcoll | step_name >> beam.ParDo(
                  _SubgraphLayerDoFn(), spec, remote_op_name
              )
  
          # Construct Beam subgraph for a remote op.
          else:
              # ExecutionSpec stores one remote op.
              child_remote_op_name = list(spec.output_names)[0]
              step_descriptor = "[Parent_%s][Child_%s]" % (
                  remote_op_name,
                  child_remote_op_name,
              )
  
              step_name = "LoadRemoteGraphInputs%s" % step_descriptor
              pcoll = (
                  pcoll
                  | step_name
                  >> _LoadRemoteGraphInputs(  # pylint: disable=no-value-for-parameter
                      remote_op_name,
                      child_remote_op_name,
                      remote_op_name_to_graph_name,
                      graph_to_remote_op_input_name_mapping,
                  )
              )
  
              # A good place to add beam.Reshuffle() to prevent fusion.
              step_name = "ExecuteGraph%s" % step_descriptor
              pcoll = (
                  pcoll
                  | step_name
                  >> ExecuteGraph(  # pylint: disable=no-value-for-parameter
                      child_remote_op_name,
                      remote_op_name_to_graph_name,
                      graph_name_to_specs,
                      graph_to_remote_op_input_name_mapping,
                  )
              )
  
              step_name = "ExtractRemoteGraphOutput%s" % step_descriptor
              pcoll = (
                  pcoll
                  | step_name
                  >> _ExtractRemoteGraphOutput(  # pylint: disable=no-value-for-parameter
                      remote_op_name,
                      child_remote_op_name,
                      remote_op_name_to_graph_name,
                      graph_name_to_specs,
                  )
              )
  
      return pcoll
  
file_path: tfx/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline.py
imports: import copy
  from typing import Any, Dict, Iterator, List, Mapping
  import apache_beam
  import tensorflow
  from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import execution_spec
label_desc: Use image node_id 1 for calling a global function with example usage: ExecuteGraph(pcoll, remote_op_name, remote_op_name_to_graph_name, graph_name_to_specs, graph_to_remote_op_input_name_mapping) and returns: pcoll