Function revert_delay_pattern

zonos/codebook_pattern.py:10–12

(codes: torch.Tensor)

def revert_delay_pattern(codes: torch.Tensor):
    _, n_q, seq_len = codes.shape
    return torch.stack([codes[:, k, k + 1 : seq_len - n_q + k + 1] for k in range(n_q)], dim=1)