code: def get_schedule_distribution(schedule, global_step=None):
    
    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])])
docs: 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.