Function shifted_unwindow

k_diffusion/models/modules.py:293–296

(window_shift, x)

def shifted_unwindow(window_shift, x):
    x = unwindow(x)
    x = torch.roll(x, shifts=(-window_shift, -window_shift), dims=(-2, -3))
    return x


@lru_cache