import torch


def init_shift1d_nfold(shift, nfold=8, noise=1e-3):
    """
    Mimic TSM hard-coded shift.
    First 1/nfold channels shift one time step to the past
    Second 1/nfold channels shift one time step to the future
    Remaining channels remain unchanged
    """
    dim, channels = shift.size()
    assert dim == 1, "only works with rubiks1d"
    with torch.no_grad():
        group = channels // nfold
        shift[:, :group] = 1
        shift[:, group : 2 * group] = -1
        # perturb to stay away from zero
        shift[:, 2 * group :].uniform_(-noise, noise)
    return shift