import torch


def get_transformer_paper_schedule(optimizer, num_warmup_steps, last_epoch=-1):
    """
    Create a schedule with a learning rate that follows from the paper "Attention is all you need",
    a linear warmup period followed by decrease that's inversely proportionl to square root of step number.

    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / max(1.0, float(num_warmup_steps)**1.5)
        return 1.0 / max(1.0, float(current_step)**0.5)
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)