import torch


def cumavg(a: torch.FloatTensor):
    """

    Parameters
    ----------
    a : torch.tensor
        Shape: l,b,e

    Returns
    -------
    Cumulative average : torch.tensor
        Shape : l,b,e
        avg over l dimension
    """
    # l,b,e
    cumsum = torch.cumsum(a, 0)
    batch_size = cumsum.shape[1]

    # Note that for the shorter sequences the cum avg is not correct
    # after the last time step of that sequence.
    # This is corrected in the loss calculation.
    lengths = torch.arange(1, 1 + cumsum.shape[0], device=a.device)

    return cumsum / lengths.reshape(cumsum.shape[0], 1, 1)