code: def testReturnsTopoSort(self, scheduler_alg):
      mtf_graph = mtf.Graph()
      mesh = mtf.Mesh(mtf_graph, "my_mesh")
      x = mtf.Constant(
          mesh,
          0,
          shape=mtf.convert_to_shape("a:3,b:4"),
          dtype=tf.int32,
          name="X",
      ).outputs[0]
      y = mtf.Constant(
          mesh,
          0,
          shape=mtf.convert_to_shape("b:4,c:5"),
          dtype=tf.int32,
          name="Y",
      ).outputs[0]
      mtf.EinsumOperation(
          [x, y], mtf.convert_to_shape("a:3,c:5"), name="Z1"
      )
      mtf.EinsumOperation(
          [x, y], mtf.convert_to_shape("a:3,c:5"), name="Z2"
      )
  
      graph = graph_interface.GraphInterface(mtf_graph)
      graph.set_tensor_final("Z1:0")
      graph.set_tensor_final("Z2:0")
      schedule = list(
          scheduler.minimize_peak_memory(graph, scheduler_alg)
      )
  
      self.assertCountEqual(schedule[0:2], [0, 1])
      self.assertCountEqual(schedule[2:4], [2, 3])
  
file_path: mesh/mesh_tensorflow/auto_mtf/scheduler_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import mesh_tensorflow
  from mesh_tensorflow.auto_mtf import graph_interface
  from mesh_tensorflow.auto_mtf import scheduler
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling the SchedulerTest obj's underlying member method code with example usage: obj.testReturnsTopoSort(scheduler_alg) without return types