jax_sgmc.scheduler.polynomial_step_size_first_last

jax_sgmc.scheduler.polynomial_step_size_first_last(first=1.0, last=1.0, gamma=0.33)

Initializes polynomial step size schedule via first and last step.

Parameters:

• first ([<class 'float'>, <class 'jax.Array'>]) – Step size in the first iteration

• last ([<class 'float'>, <class 'jax.Array'>]) – Step size in the last iteration

• gamma ([<class 'float'>, <class 'jax.Array'>]) – Rate of decay

Return type:

specific_scheduler

Returns:

Returns a polynomial step size schedule defined via the first and the last step size.