import torch


def periodic_cumsum(tensor, period, dim):
    """
    Compute periodic cumsum along a given dimension. For example if dim=0::

        for t in range(period):
            assert result[t] == tensor[t]
        for t in range(period, len(tensor)):
            assert result[t] == tensor[t] + result[t - period]

    This is useful for computing drifting seasonality in time series models.

    :param torch.Tensor tensor: A tensor of differences.
    :param int period: The period of repetition.
    :param int dim: The tensor dimension along which to accumulate.
    """
    assert isinstance(period, int) and period > 0
    assert isinstance(dim, int)
    if dim >= 0:
        dim -= tensor.dim()

    # Pad to even size.
    size = tensor.size(dim)
    repeats = (size + period - 1) // period
    padding = repeats * period - size
    if torch._C._get_tracing_state() or padding:
        tensor = torch.nn.functional.pad(tensor, (0, 0) * (-1 - dim) + (0, padding))

    # Accumulate.
    shape = tensor.shape[:dim] + (repeats, period) + tensor.shape[tensor.dim() + dim + 1:]
    result = tensor.reshape(shape).cumsum(dim=dim - 1).reshape(tensor.shape)

    # Truncate to original size.
    if torch._C._get_tracing_state() or padding:
        result = result[(Ellipsis, slice(None, size)) + (slice(None),) * (-1 - dim)]
    return result