ott.tools.progot.get_alpha_schedule

ott.tools.progot.get_alpha_schedule#

ott.tools.progot.get_alpha_schedule(kind, *, num_steps)[source]#

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