ott.tools.progot.get_alpha_schedule

ott.tools.progot.get_alpha_schedule(kind, *, num_steps)

Get the step size schedule.

Convenience wrapper to get a sequence of num_steps timestamps between 0 and 1, distributed according to the kind option below. See Section 4 in [Kassraie et al., 2024] for more details.

Parameters:

• kind (Literal['lin', 'exp', 'quad']) –

  The schedule to create:

  • 'lin' - constant-speed schedule.

  • 'exp' - decelerating schedule.

  • 'quad' - accelerating schedule.

• num_steps (int) – Total number of steps.

Return type:

Array

Returns:

The stepsize schedule, array of shape [num_steps,].