def rank_correction(measure, n, rank=1, dtype=torch.float):
    """
    Return low-rank matrix L such that A + L is normal
    Args:
        measure (str): the type of measure
            legt - Legendre (translated)
            legs - Legendre (scaled)
            glagt - generalized Laguerre (translated)
            lagt, tlagt - previous versions of (tilted) Laguerre with slightly different normalization
        n (int): dimension
        rank (int): rank
        dtype (str): data type
    """

    if measure == 'legs':
        assert rank >= 1
        p = torch.sqrt(.5+torch.arange(n, dtype=dtype)).unsqueeze(0)  # (1 n)
    elif measure == 'legt':
        assert rank >= 2
        p = torch.sqrt(1+2*torch.arange(n, dtype=dtype))  # (n)
        p0 = p.clone()
        p0[0::2] = 0.
        p1 = p.clone()
        p1[1::2] = 0.
        p = torch.stack([p0, p1], dim=0)  # (2 n)
    elif measure == 'lagt':
        assert rank >= 1
        p = .5**.5 * torch.ones(1, n, dtype=dtype)
    else:
        raise NotImplementedError

    d = p.size(0)
    if rank > d:
        p = torch.stack([p, torch.zeros(n, dtype=dtype).repeat(rank-d, d)], dim=0)  # (rank n)
    return p
Return low-rank matrix L such that A + L is normal
Args:
measure (str): the type of measure
legt - Legendre (translated)
legs - Legendre (scaled)
glagt - generalized Laguerre (translated)
lagt, tlagt - previous versions of (tilted) Laguerre with slightly different normalization
n (int): dimension
rank (int): rank
dtype (str): data type
towhee/models/vis4mer/utils.py