# Path: mesh_tensorflow/auto_mtf/graph_interface.py
# class GraphInterface(object):
#   def __init__(self, graph, canonical_device=None):
#   def get_num_operations(self):
#   def get_all_operation_names(self):
#   def get_operation_input_names(self, operation_name):
#   def get_operation_output_names(self, operation_name):
#   def get_all_tensor_names(self):
#   def get_tensor_dtype(self, tensor_name):
#   def get_tensor_shape(self, tensor_name):
#   def get_tensor_num_entries(self, tensor_name, partial_layout=None,
#                              mesh_dimension_to_size=None):
#   def get_tensor_size(self, tensor_name, partial_layout=None,
#                       mesh_dimension_to_size=None):
#   def get_tensor_device(self, tensor_name):
#   def is_tensor_on_canonical_device(self, tensor_name):
#   def get_operation_device(self, operation_name):
#   def get_tensor_mtf_dimension_names(self, tensor_name):
#   def get_operation_mtf_dimension_names(self, operation_name):
#   def set_tensor_final(self, tensor_name):
#   def is_tensor_final(self, tensor_name):
#   def compute_cost_graph(self, devices=None):
#   def compute_memory_contents_under_schedule(self, schedule):
#   def _initialize_operations(self):
#   def _initialize_operation_name_to_id(self):
#   def _initialize_tensor_name_to_ids(self):
#   def _get_tensors(self):
#   def _name_to_operation(self, operation_name):
#   def _name_to_tensor(self, tensor_name):
#   def _compute_initial_out_degree(self):


import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import graph_interface
from tensorflow.core.framework import cost_graph_pb2
from tensorflow.core.framework import tensor_shape_pb2
from tensorflow.core.framework import types_pb2
    self.assertCountEqual(graph.get_operation_mtf_dimension_names("Z1"),
                          ["a", "b", "c"])
    self.assertCountEqual(graph.get_operation_mtf_dimension_names("Z2"),
                          ["a", "b", "c"])

    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("X:0"),
                          ["a", "b"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Y:0"),
                          ["b", "c"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Z1:0"),
                          ["a", "c"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Z1:0"),
                          ["a", "c"])

    self.assertIsNone(graph.get_tensor_device("X:0"))
    self.assertIsNone(graph.get_tensor_device("Y:0"))
    self.assertIsNone(graph.get_tensor_device("Z1:0"))
    self.assertIsNone(graph.get_tensor_device("Z2:0"))

    self.assertTrue(graph.is_tensor_on_canonical_device("X:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Y:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Z1:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Z2:0"))

    self.assertEqual(graph.compute_cost_graph().SerializeToString(),
                     self._deviceless_cost_graph_string)
    self.assertEqual(graph.compute_cost_graph(devices=[]).SerializeToString(),
                     self._deviceless_cost_graph_string)

  def testNotAGraph(self):

# Path: mesh_tensorflow/auto_mtf/graph_interface.py
# class GraphInterface(object):
#   def __init__(self, graph, canonical_device=None):
#   def get_num_operations(self):
#   def get_all_operation_names(self):
#   def get_operation_input_names(self, operation_name):
#   def get_operation_output_names(self, operation_name):
#   def get_all_tensor_names(self):
#   def get_tensor_dtype(self, tensor_name):
#   def get_tensor_shape(self, tensor_name):
#   def get_tensor_num_entries(self, tensor_name, partial_layout=None,
#                              mesh_dimension_to_size=None):
#   def get_tensor_size(self, tensor_name, partial_layout=None,
#                       mesh_dimension_to_size=None):
#   def get_tensor_device(self, tensor_name):
#   def is_tensor_on_canonical_device(self, tensor_name):
#   def get_operation_device(self, operation_name):
#   def get_tensor_mtf_dimension_names(self, tensor_name):
#   def get_operation_mtf_dimension_names(self, operation_name):
#   def set_tensor_final(self, tensor_name):
#   def is_tensor_final(self, tensor_name):
#   def compute_cost_graph(self, devices=None):
#   def compute_memory_contents_under_schedule(self, schedule):
#   def _initialize_operations(self):
#   def _initialize_operation_name_to_id(self):
#   def _initialize_tensor_name_to_ids(self):
#   def _get_tensors(self):
#   def _name_to_operation(self, operation_name):
#   def _name_to_tensor(self, tensor_name):
#   def _compute_initial_out_degree(self):


# Path: mesh_tensorflow/auto_mtf/graph_interface_test.py
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import graph_interface
from tensorflow.core.framework import cost_graph_pb2
from tensorflow.core.framework import tensor_shape_pb2
from tensorflow.core.framework import types_pb2

    self.assertCountEqual(graph.get_operation_mtf_dimension_names("Z1"),
                          ["a", "b", "c"])
    self.assertCountEqual(graph.get_operation_mtf_dimension_names("Z2"),
                          ["a", "b", "c"])

    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("X:0"),
                          ["a", "b"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Y:0"),
                          ["b", "c"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Z1:0"),
                          ["a", "c"])
    self.assertCountEqual(graph.get_tensor_mtf_dimension_names("Z1:0"),
                          ["a", "c"])

    self.assertIsNone(graph.get_tensor_device("X:0"))
    self.assertIsNone(graph.get_tensor_device("Y:0"))
    self.assertIsNone(graph.get_tensor_device("Z1:0"))
    self.assertIsNone(graph.get_tensor_device("Z2:0"))

    self.assertTrue(graph.is_tensor_on_canonical_device("X:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Y:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Z1:0"))
    self.assertTrue(graph.is_tensor_on_canonical_device("Z2:0"))

    self.assertEqual(graph.compute_cost_graph().SerializeToString(),
                     self._deviceless_cost_graph_string)
    self.assertEqual(graph.compute_cost_graph(devices=[]).SerializeToString(),
                     self._deviceless_cost_graph_string)

  def testNotAGraph(self):

    self.assertRaises(TypeError, graph_interface.GraphInterface, "hello")