code: def get_schedule_distribution(schedule, global_step=None):
      """Computes the pmf of a schedule given the global_step.
  
      Args:
        schedule: A schedule tuple, see encode_schedule for details.
        global_step: A scalar tensor, the step to query the schedule.
  
      Returns:
        A 1-D tensor of probs, the sampling distribution of the global_step.
      """
      interpolation, steps, pmfs = schedule
      if len(pmfs) == 1:
          # py_func doesn't seem to work on TPU - at least get the constant case to
          # run.
          # TODO(noam): get the general case working.
          return pmfs[0]
      if global_step is None:
          global_step = tf.train.get_or_create_global_step()
      if interpolation == "step":
          interpolation_fn = step_interpolation
      elif interpolation == "linear":
          interpolation_fn = linear_interpolation
      else:
          raise ValueError(
              "Invalid interpolation strategy: %s" % interpolation
          )
      return tf.reshape(
          tf.py_func(
              func=lambda x: interpolation_fn(
                  x, np.array(steps), np.array(pmfs)
              ),
              inp=[global_step],
              Tout=tf.float32,
          ),
          [len(pmfs[0])],
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import numpy
  from tensor2tensor.data_generators import problem
  from tensor2tensor.data_generators import text_problems
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: get_schedule_distribution(schedule, global_step) and returns: tf, pmfs