class_docstr: A test for the beam pipeline library.
docstr: Compares the results from the original model and the beam pipeline.
docstr_tok: ["Compares","the","results","from","the","original","model","and","the","beam","pipeline","."]
code: def test_compare_outputs(self):
      """Compares the results from the original model and the beam pipeline."""
      with tempfile.TemporaryDirectory() as temp_dir:
          create_complex_graph.save_examples_as_graphdefs(temp_dir)
  
          # "main" is both a graph name and a remote op name.
          root_graph = "main"
          graph_to_remote_op_input_name_mapping = {
              "main": {
                  "remote_op_a": {"ids_a": "ids1"},
                  "remote_op_b": {"ids_b1": "ids1", "ids_b2": "ids2"},
                  "remote_op_a_1": {"ids_a": "FloorMod_1"},
                  "remote_op_b_1": {
                      "ids_b1": "FloorMod_1",
                      "ids_b2": "FloorMod",
                  },
              },
              "graph_b": {
                  "remote_op_a": {"ids_a": "FloorMod"},
                  "remote_op_a_1": {"ids_a": "ids_b2"},
              },
          }
  
          # Create input PColl with this.
          root_graph_input_data = [
              {"main": {"import/ids1:0": 3, "import/ids2:0": 3}},
              {"main": {"import/ids1:0": 10, "import/ids2:0": 10}},
          ]
  
          graph_name_to_filepath = {
              "main": os.path.join(temp_dir, "main_graph.pb"),
              "graph_b": os.path.join(temp_dir, "graph_b.pb"),
              "graph_a": os.path.join(temp_dir, "graph_a.pb"),
          }
          graph_name_to_output_names = {
              "main": ["AddN_1"],
              "graph_b": ["Add_1"],
              "graph_a": ["embedding_lookup/Identity"],
          }
          remote_op_name_to_graph_name = {
              "main": "main",
              "remote_op_a": "graph_a",
              "remote_op_a_1": "graph_a",
              "remote_op_b": "graph_b",
              "remote_op_b_1": "graph_b",
          }
  
          original_model_outputs = _run_original_model(
              root_graph,
              root_graph_input_data,
              graph_name_to_filepath,
              graph_name_to_output_names,
          )
  
          graph_name_to_graph_def = (
              graph_partition.get_graph_name_to_graph_def(
                  graph_name_to_filepath
              )
          )
          graph_name_to_specs = graph_partition.partition_all_graphs(
              graph_name_to_graph_def, graph_name_to_output_names
          )
  
          with beam.Pipeline() as p:
              beam_inputs = p | "LoadInputs" >> beam.Create(
                  root_graph_input_data
              )
              beam_outputs = (
                  beam_inputs
                  | "RunModel"
                  >> beam_pipeline.ExecuteGraph(
                      root_graph,
                      remote_op_name_to_graph_name,
                      graph_name_to_specs,
                      graph_to_remote_op_input_name_mapping,
                  )
                  | "ExtractOutputs"
                  >> beam.Map(
                      _extract_outputs,
                      root_graph,
                      graph_name_to_output_names,
                  )
              )
  
              util.assert_that(
                  beam_outputs, almost_equal_to(original_model_outputs)
              )
  
file_path: tfx/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py
imports: import os
  import tempfile
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import beam_pipeline
  from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import create_complex_graph
  from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import graph_partition
label_desc: Use image node_id 1 for calling the BeamPipelineTest obj's underlying member method code with example usage: obj.test_compare_outputs() without return types