Scale updates using a custom schedule for the step size.
JIT-compatible when the schedule body uses np.* array operations.
np.*
Scale updates using a custom schedule for the step size.
JIT-compatible when the schedule body uses
np.*array operations.