def learning_rate(self, step):
    """Return the learning rate for the given step."""
    if self._lr_schedule is not None:
      with fastmath.use_backend(fastmath.Backend.NUMPY):
        return self._lr_schedule(step)
    opt = self._optimizer
    if callable(opt):  # when optimizer is a function, like Adam, not Adam()
      opt = opt()
    params = opt._init_opt_params  # pylint: disable=protected-access
    return params['learning_rate']
Return the learning rate for the given step.
trax/supervised/training.py