def build_nonbipartite_schedule_inner()

in torchbiggraph/train_gpu.py [0:0]


def build_nonbipartite_schedule_inner(size: int) -> List[List[int]]:
    if size <= 0 or size % 2 != 0:
        raise ValueError("Bad")
    if size == 2:
        return [[(0, 1), (1, 0)]]
    half = size // 2
    pre = [[(i, (i + j) % half + half) for j in range(half)] for i in range(half)]
    post = [[(i + half, (i + j) % half) for j in range(half)] for i in range(half)]
    mid = build_nonbipartite_schedule_inner(half)
    res = []
    res.extend([pre[i] + mid[i] + post[i] for i in range(half // 2)])
    res.extend(
        [
            pre[i + half // 2]
            + [(x + half, y + half) for x, y in mid[i]]
            + post[i + half // 2]
            for i in range(half // 2)
        ]
    )
    return res