def get_drop_path_keep_prob(keep_prob_for_last_stage, schedule,
                            current_stage, num_stages):
  """Gets drop path keep probability for current stage.

  Args:
    keep_prob_for_last_stage: A float, the drop path keep probability for
      last stage. This flag is used in conjunction with the flag `schedule`, as
      they together determine drop path keep probability for the other stages.
    schedule: A string, the drop path schedule. Currently, we support
      'constant': use the same drop path keep probability for all stages, and
      'linear': linearly decrease the drop path keep probability from 1.0 at
        0-th stage (or STEM) to `keep_prob_for_last_stage` at last stage.
    current_stage:  An integer, current stage number.
    num_stages: An integer, the number of stages.

  Returns:
    The drop path keep probability for the current stage.

  Raises:
    ValueError: If schedule is not supported.
  """
  if schedule == 'constant':
    return keep_prob_for_last_stage
  elif schedule == 'linear':
    return 1.0 - (1.0 - keep_prob_for_last_stage) * current_stage / num_stages
  else:
    raise ValueError('Unexpected schedule %s.' % schedule)