def de_vec(a, rows, cols):
    """Column-major de-vectorization. First dimension is batch dim.
    """
    batch = a.size(0)
    # dim = math.floor(math.sqrt(a.size(1)))
    a_mat = a.reshape(batch, cols, rows).permute(0, 2, 1)
    return a_mat