code: def test_get_schedule_distribution(self, schedule, steps, pmfs):
      with self.test_session() as sess:
          global_step = tf.train.get_or_create_global_step()
          output = multi_problem_v2.get_schedule_distribution(
              schedule, global_step
          )
          sess.run(global_step.initializer)
          for step, pmf in zip(steps, pmfs):
              sess.run(global_step.assign(step))
              self.assertAllClose(sess.run(output), pmf)
  
file_path: tensor2tensor/tensor2tensor/data_generators/multi_problem_v2_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  from tensor2tensor.data_generators import multi_problem_v2
  from tensor2tensor.data_generators import problem
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the MultiProblemV2Test obj's underlying member method code with example usage: obj.test_get_schedule_distribution(schedule, steps, pmfs) without return types