import torch


def lin_momentum(mu_prev, mu_curr, mu_stream,
                 momentum, momentum_pow, momentum_batch):
    """
    Helper function for performing an exponential moving average, ema, using
    convolutions which can be distributed across compute fabric.

    Arguments:
        mu_prev: previous time steps statistic
        mu_curr: this time steps statistic
        mu_stream: ema from last time step
        momentum: 1 - 1 / h
        momentum_pow: momentum ^ range(b_size - 1, -1, -1)
        momentum_batch: momentum ^ b_size of size :math:`(N, L)`

    Return:
        updated mu_stream (stale): to use for fprop
        updated mu_stream (current): to cache for next iteration
    """
    input = torch.cat([mu_prev[1:], mu_curr]).transpose(0, 1).unsqueeze(1)
    tmp = torch.nn.functional.conv1d(input, momentum_pow).squeeze().transpose(0, 1)
    curr = (momentum_batch * mu_stream + (1 - momentum) * tmp)

    return torch.cat([mu_stream[-1].unsqueeze(0), curr[:-1]]), curr