class_docstr: A set of tests for the graph partitioning library.
class_docstr_tok: ["A","set","of","tests","for","the","graph","partitioning","library","."]
docstr: Sets up some example graphs and their partitions.
docstr_tok: ["Sets","up","some","example","graphs","and","their","partitions","."]
code: def setUp(self):
      """Sets up some example graphs and their partitions."""
      super().setUp()
      with tempfile.TemporaryDirectory() as temp_dir:
          # Save examples into a temporary directory
          create_complex_graph.save_examples_as_graphdefs(temp_dir)
  
          graph_name_to_filepath = {
              "main": os.path.join(temp_dir, "main_graph.pb"),
              "remote_op_a": os.path.join(temp_dir, "graph_a.pb"),
              "remote_op_b": os.path.join(temp_dir, "graph_b.pb"),
          }
          graph_name_to_outputs = {
              "main": ["AddN_1"],
              "remote_op_b": ["Add_1"],
              "remote_op_a": ["embedding_lookup/Identity"],
          }
  
          graph_name_to_graph_def = (
              graph_partition.get_graph_name_to_graph_def(
                  graph_name_to_filepath
              )
          )
          self.graph_name_to_specs = (
              graph_partition.partition_all_graphs(
                  graph_name_to_graph_def, graph_name_to_outputs
              )
          )
  
file_path: tfx/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py
imports: import os
  import tempfile
  import tensorflow
  from tfx.dsl.io import fileio
  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
  from google.protobuf import text_format
label_desc: Use image node_id 1 for calling the PartitionTest obj's underlying member method code with example usage: obj.setUp() without return types