Returns an optimizer with scheduled learning rate.

The returned optimizer will use a learning rate of schedule_fn(i) for the i-th invocation of its next method, indexing from 0.

optimizer A tff.learning.optimizers.Optimizer which uses a learning rate.
schedule_fn A callable mapping integer round number to a floating point learning rate. To be invoked in the context of a tff.tensorflow.computation, thus should support a tf.Tensor input.

A tff.learning.optimizers.Optimizer.

KeyError If the provided optimizer's state is not a dictionary with learning rate stored under the tff.learning.optimizers.LEARNING_RATE_KEY key.