optax.scale_by_schedule

optax.scale_by_schedule(step_size_fn: base.Schedule) → base.GradientTransformation

Scale updates using a custom schedule for the step_size.

Parameters:

step_size_fn – A function that takes an update count as input and proposes the step_size to multiply the updates by.

Returns:

A optax.GradientTransformation object.