code: def _ExtractRemoteGraphOutput(  # pylint: disable=invalid-name
      pcoll: beam.pvalue.PCollection,
      parent_remote_op_name: str,
      child_remote_op_name: str,
      remote_op_name_to_graph_name: Mapping[str, str],
      graph_name_to_specs: Mapping[
          str, List[execution_spec.ExecutionSpec]
      ],
  ) -> beam.pvalue.PCollection:
      """A PTransform that extracts remote graph output.
  
      After finish executing a remote graph, we need to collect its output.
      We first find the output name of the remote graph, then we copy the
      output of the remote graph to its parent graph. Finally, we clear the
      intermediate results of the remote graph.
  
      Note we assumed that each node has only one output, which also applies
      to remote op. This means that a remote graph can only have one output.
  
      Args:
        pcoll: A PCollection of child graph results. Each element is a dictionary
          from remote op names to a dictionary from tensor names to values. Here,
          element[child_remote_op_name] stores graph inputs, intermediate results,
          and graph output.
        parent_remote_op_name: The remote op name of the parent graph.
        child_remote_op_name: The remote op name of the child 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.
  
      Returns:
        A PCollection of child graph output in parent graph. Each element is a
        dictionary from remote op names to a dictionary from tensor names to
        values. Here, element[parent_remote_op_name] contains the output from
        the child graph, and element[child_remote_op_name] is deleted.
      """
      child_graph_name = remote_op_name_to_graph_name[
          child_remote_op_name
      ]
      child_specs = graph_name_to_specs[child_graph_name]
      child_output_name = list(child_specs[-1].output_names)[0]
  
      step_name_extract = "ExtractOutput[Graph_%s][Output_%s]" % (
          child_remote_op_name,
          child_output_name,
      )
      step_name_clear = "ClearIntermediateOutputs[Graph_%s]" % (
          child_remote_op_name
      )
  
      return (
          pcoll
          | step_name_extract
          >> beam.Map(
              _copy_tensor_value,
              child_remote_op_name,
              _import_tensor_name(child_output_name),
              parent_remote_op_name,
              _import_tensor_name(child_remote_op_name),
          )
          | step_name_clear
          >> beam.Map(
              _clear_outputs_for_finished_graph, child_remote_op_name
          )
      )
  
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 5 for calling a global function with example usage: _ExtractRemoteGraphOutput(pcoll, parent_remote_op_name, child_remote_op_name, remote_op_name_to_graph_name, graph_name_to_specs) and returns: unknown