# 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/scheduler.py
# def minimize_peak_memory(graph, scheduler_alg):
# def _minimize_peak_memory_naive(graph):
# def _minimize_peak_memory_list(graph):


from absl.testing import parameterized
from mesh_tensorflow.auto_mtf import graph_interface
from mesh_tensorflow.auto_mtf import scheduler
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
    mesh = mtf.Mesh(mtf_graph, 'my_mesh')
    mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('a:4'), dtype=tf.int32,
                 name='X')
    y = mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('b:3'), dtype=tf.int32,
                     name='Y').outputs[0]
    mtf.BroadcastOperation(y, mtf.convert_to_shape('b:3,c:2'), name='Z')

    graph = graph_interface.GraphInterface(mtf_graph)
    graph.set_tensor_final('X:0')
    graph.set_tensor_final('Z:0')
    schedule = list(scheduler.minimize_peak_memory(graph, 'LIST'))
    # When nothing is scheduled:
    #   X frees -4 entries
    #   Y frees -3 entries
    # After [Y] scheduled:
    #   X frees -4 entries
    #   Z frees -3 entries
    # Hence the schedule should be [Y, Z, X].
    self.assertEqual(schedule, [1, 2, 0])

  def testMinimizePeakMemoryList_ZeroUseTensor(self):
    mtf_graph = mtf.Graph()
    mesh = mtf.Mesh(mtf_graph, 'my_mesh')
    mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('a:4'), dtype=tf.int32,
                 name='X')
    y = mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('b:3'), dtype=tf.int32,
                     name='Y').outputs[0]
    mtf.BroadcastOperation(y, mtf.convert_to_shape('b:3,c:2'), name='Z')

    graph = graph_interface.GraphInterface(mtf_graph)

# 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/scheduler.py
# def minimize_peak_memory(graph, scheduler_alg):
# def _minimize_peak_memory_naive(graph):
# def _minimize_peak_memory_list(graph):


# Path: mesh_tensorflow/auto_mtf/scheduler_test.py
from absl.testing import parameterized
from mesh_tensorflow.auto_mtf import graph_interface
from mesh_tensorflow.auto_mtf import scheduler
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf

    mesh = mtf.Mesh(mtf_graph, 'my_mesh')
    mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('a:4'), dtype=tf.int32,
                 name='X')
    y = mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('b:3'), dtype=tf.int32,
                     name='Y').outputs[0]
    mtf.BroadcastOperation(y, mtf.convert_to_shape('b:3,c:2'), name='Z')

    graph = graph_interface.GraphInterface(mtf_graph)
    graph.set_tensor_final('X:0')
    graph.set_tensor_final('Z:0')
    schedule = list(scheduler.minimize_peak_memory(graph, 'LIST'))
    # When nothing is scheduled:
    #   X frees -4 entries
    #   Y frees -3 entries
    # After [Y] scheduled:
    #   X frees -4 entries
    #   Z frees -3 entries
    # Hence the schedule should be [Y, Z, X].
    self.assertEqual(schedule, [1, 2, 0])

  def testMinimizePeakMemoryList_ZeroUseTensor(self):
    mtf_graph = mtf.Graph()
    mesh = mtf.Mesh(mtf_graph, 'my_mesh')
    mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('a:4'), dtype=tf.int32,
                 name='X')
    y = mtf.Constant(mesh, 0, shape=mtf.convert_to_shape('b:3'), dtype=tf.int32,
                     name='Y').outputs[0]
    mtf.BroadcastOperation(y, mtf.convert_to_shape('b:3,c:2'), name='Z')

    graph = graph_interface.GraphInterface(mtf_graph)

    schedule = list(scheduler.minimize_peak_memory(graph, 'LIST'))