import torch


def discounted_cumsum(rewards, discount):

    """Calculates the cummulative sum of discounted rewards
    Arguments:
        rewards {torch.Tensor} -- rewards
        discount {float} -- discount factor
    Returns:
        [type] -- cummulative sum of discounted rewards
    """
    discount **= torch.arange(0, rewards.shape[0])
    disc_cumsum = torch.cumsum(discount * rewards, 0).flip()
    return disc_cumsum