def remap_jigsaw()

in train.py [0:0]


def remap_jigsaw(x, order):
    r, c, ch = x.shape
    g = H.jigsaw_grid_size
    gr, gc = r // g, c // g
    x = x.reshape((g, gr, g, gc, ch))
    x = np.transpose(x, [0, 2, 1, 3, 4])
    x = x.reshape([g * g, gr, gc, ch])
    perm = H.jigsaw_perms[order]
    x = x[perm, :, :, :]
    x = x.reshape([g, g, gr, gc, ch])
    x = np.transpose(x, [0, 2, 1, 3, 4])
    x = x.reshape((r, c, ch))
    return x