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])])
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.