jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _fit_2d(self, data: Tensor) -> None:
        _, columns = data.shape
        mean = torch.mean(data, dim=0)
        eigenvalues, eigenvectors = [], []
        for column in range(columns):
            mean_ = mean[column]
            data_column = data[:, column]
            shifted_data = data_column - mean_
            cov = torch.dot(shifted_data.t(), shifted_data)
            eigenvalue, eigenvector = torch.symeig(cov.view(1, 1), eigenvectors=True)
            eigenvalues.append(eigenvalue)
            eigenvectors.append(eigenvector)

        self.params = {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}
//  [python]
def _fit_2d(self, data: Tensor) -> None:
        _, columns = data.shape
        mean = torch.mean(data, dim=0)
        eigenvalues, eigenvectors = [], []
        for column in range(columns):
            mean_ = mean[column]
            data_column = data[:, column]
            shifted_data = data_column - mean_
            cov = torch.dot(shifted_data.t(), shifted_data)
            eigenvalue, eigenvector = torch.linalg.eigh(cov.view(1, 1))
            eigenvalues.append(eigenvalue)
            eigenvectors.append(eigenvector)

        self.params = {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}
//  [python]
def eig_decomp_reconstruct(w):
    n_output, n_input, kh, hw = w.size()
    assert n_output == n_input
    n = n_output
    # make a list of tensors of shape NxN
    tensors = map(torch.squeeze, w.view(n, n, -1).split(1, dim=2))

    def evd(x):
        e, V = torch.symeig(x.data, eigenvectors=True) # only use triu(x)
        x_hat = torch.mm(torch.mm(V, torch.diag(e)), V.t())
        return Variable(x_hat)

    # apply hat(u) = V diag(e) V' to every tensor u and stack back
    w_hat = torch.cat([evd(u).unsqueeze(2) for u in tensors], dim=2)
    return w_hat.view_as(w)
//  [python]
def eig_decomp_reconstruct(w):
    n_output, n_input, kh, hw = w.size()
    assert n_output == n_input
    n = n_output
    # make a list of tensors of shape NxN
    tensors = map(torch.squeeze, w.view(n, n, -1).split(1, dim=2))

    def evd(x):
        e, V = torch.linalg.eigh(x.data) # only use triu(x)
        x_hat = torch.mm(torch.mm(V, torch.diag(e)), V.t())
        return Variable(x_hat)

    # apply hat(u) = V diag(e) V' to every tensor u and stack back
    w_hat = torch.cat([evd(u).unsqueeze(2) for u in tensors], dim=2)
    return w_hat.view_as(w)
//  [python]
def refparams_vecgsm(self, org):
        ssarr, l_arr, cu_arr = [], [], []
        for i in range(len(self.subbands)):
            sub=self.subbands[i]-1
            y=org[sub]
            M = self.M
            newsizeX=int(np.floor(y.shape[2]/M)*M)
            newsizeY=int(np.floor(y.shape[3]/M)*M)
            y=y[:,:,:newsizeX,:newsizeY]
            B,C,H,W = y.shape
                        
            temp=[]
            for j in range(M):
                for k in range(M):
                    temp.append(y[:,:,k:H-(M-k)+1, j:W-(M-j)+1].reshape(B,C,-1))
            temp = torch.stack(temp,dim=3)
            mcu = torch.mean(temp,dim=2).unsqueeze(2).repeat(1,1,temp.shape[2],1)
            cu=torch.matmul((temp-mcu).permute(0,1,3,2),temp-mcu)/temp.shape[2]

            temp=[]
            for j in range(M):
                for k in range(M):
                    temp.append(y[:,:,k:H+1:M, j:W+1:M].reshape(B,C,-1))
            temp = torch.stack(temp,dim=2)
            ss=torch.matmul(torch.pinverse(cu),temp)
            # ss = torch.matmul(torch.pinverse(cu),temp)
            ss=torch.sum(ss*temp,dim=2)/(M*M)
            ss=ss.reshape(B,C,H//M,W//M)    
            v,_ = torch.symeig(cu,eigenvectors=True)
            l_arr.append(v)
            ssarr.append(ss)
            cu_arr.append(cu)

        return ssarr, l_arr, cu_arr
//  [python]
def refparams_vecgsm(self, org):
        ssarr, l_arr, cu_arr = [], [], []
        for i in range(len(self.subbands)):
            sub=self.subbands[i]-1
            y=org[sub]
            M = self.M
            newsizeX=int(np.floor(y.shape[2]/M)*M)
            newsizeY=int(np.floor(y.shape[3]/M)*M)
            y=y[:,:,:newsizeX,:newsizeY]
            B,C,H,W = y.shape
                        
            temp=[]
            for j in range(M):
                for k in range(M):
                    temp.append(y[:,:,k:H-(M-k)+1, j:W-(M-j)+1].reshape(B,C,-1))
            temp = torch.stack(temp,dim=3)
            mcu = torch.mean(temp,dim=2).unsqueeze(2).repeat(1,1,temp.shape[2],1)
            cu=torch.matmul((temp-mcu).permute(0,1,3,2),temp-mcu)/temp.shape[2]

            temp=[]
            for j in range(M):
                for k in range(M):
                    temp.append(y[:,:,k:H+1:M, j:W+1:M].reshape(B,C,-1))
            temp = torch.stack(temp,dim=2)
            ss=torch.matmul(torch.pinverse(cu),temp)
            # ss = torch.matmul(torch.pinverse(cu),temp)
            ss=torch.sum(ss*temp,dim=2)/(M*M)
            ss=ss.reshape(B,C,H//M,W//M)    
            v,_ = torch.linalg.eigh(cu)
            l_arr.append(v)
            ssarr.append(ss)
            cu_arr.append(cu)

        return ssarr, l_arr, cu_arr
//  [python]
def main():
    from torch.autograd import gradcheck

    k = torch.randn(5, 20, 20).double()
    M = k @ k.transpose(-1,-2)

    s1 = symsqrt_v1(M, func='symeig')
    test = torch.allclose(M, s1 @ s1.transpose(-1,-2))
    print('Via symeig:', test)

    s2 = symsqrt_v1(M, func='svd')
    test = torch.allclose(M, s2 @ s2.transpose(-1,-2))
    print('Via svd:  ', test)

    print('Sqrtm with symeig and svd match:', torch.allclose(s1,s2))

    M.requires_grad = True

    ## Check gradients for symsqrt
    _sqrt = partial(symsqrt, func='svd')
    test = gradcheck(_sqrt, (M,))
    print('Grach Check for sqrtm/svd:', test)

    ## Check symeig itself
    S = torch.rand(5,20,20, requires_grad=True).double()
    def func(S):
        x = 0.5 * (S + S.transpose(-2, -1))
        return torch.symeig(x, eigenvectors=True)
    print('Grad check for symeig', gradcheck(func, [S]))

    ## Check gradients for symsqrt with symeig
    _sqrt = partial(symsqrt, func='symeig')
    test = gradcheck(_sqrt, (M,))
    print('Grach Check for sqrtm/symeig:', test)
//  [python]
def main():
    from torch.autograd import gradcheck

    k = torch.randn(5, 20, 20).double()
    M = k @ k.transpose(-1,-2)

    s1 = symsqrt_v1(M, func='symeig')
    test = torch.allclose(M, s1 @ s1.transpose(-1,-2))
    print('Via symeig:', test)

    s2 = symsqrt_v1(M, func='svd')
    test = torch.allclose(M, s2 @ s2.transpose(-1,-2))
    print('Via svd:  ', test)

    print('Sqrtm with symeig and svd match:', torch.allclose(s1,s2))

    M.requires_grad = True

    ## Check gradients for symsqrt
    _sqrt = partial(symsqrt, func='svd')
    test = gradcheck(_sqrt, (M,))
    print('Grach Check for sqrtm/svd:', test)

    ## Check symeig itself
    S = torch.rand(5,20,20, requires_grad=True).double()
    def func(S):
        x = 0.5 * (S + S.transpose(-2, -1))
        return torch.linalg.eigh(x, eigenvectors=True)
    print('Grad check for symeig', gradcheck(func, [S]))

    ## Check gradients for symsqrt with symeig
    _sqrt = partial(symsqrt, func='symeig')
    test = gradcheck(_sqrt, (M,))
    print('Grach Check for sqrtm/symeig:', test)
//  [python]
def build(self, normalize=True):

        self.normalize = normalize
        self.var /= self.N
        self.mean /= self.N
        self.var = self.var - self.mean.view(1, -1) * self.mean.view(-1, 1)
        k = self.var.size(0)

        e_vals, e_vects = torch.symeig(self.var, eigenvectors=True)
        self.PCA_mul = e_vects.t().view(1, k, k).clone()
        self.PCA_values = e_vals.clone()
        self.building = False
//  [python]
def build(self, normalize=True):

        self.normalize = normalize
        self.var /= self.N
        self.mean /= self.N
        self.var = self.var - self.mean.view(1, -1) * self.mean.view(-1, 1)
        k = self.var.size(0)

        e_vals, e_vects = torch.linalg.eigh(self.var)
        self.PCA_mul = e_vects.t().view(1, k, k).clone()
        self.PCA_values = e_vals.clone()
        self.building = False
//  [python]
def forward(ctx, input, eps=1e-2):
        use_cuda = input.is_cuda
        if input.shape[0] < 300:
            input = input.cpu()
        e, v = torch.symeig(input, eigenvectors=True)
        if use_cuda and not e.is_cuda:
            e = e.cuda()
            v = v.cuda()
        e = e.clamp(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.e_sqrt = e_sqrt
        ctx.v = v
        e_rsqrt = e_sqrt.reciprocal()

        output = v.mm(torch.diag(e_rsqrt).mm(v.t()))
        return output
//  [python]
def forward(ctx, input, eps=1e-2):
        use_cuda = input.is_cuda
        if input.shape[0] < 300:
            input = input.cpu()
        e, v = torch.linalg.eigh(input)
        if use_cuda and not e.is_cuda:
            e = e.cuda()
            v = v.cuda()
        e = e.clamp(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.e_sqrt = e_sqrt
        ctx.v = v
        e_rsqrt = e_sqrt.reciprocal()

        output = v.mm(torch.diag(e_rsqrt).mm(v.t()))
        return output
//  [python]
def latent_correlation_layer(self, x):
        batch_size = x.shape[0]
        output, hidden = self.GRU(x)
        output = output[:, -1:, :]
        attention = self.self_graph_attention(output)
        attention = torch.mean(attention, dim=0)  # average of batch
        degree = torch.sum(attention, dim=1)
        # laplacian is sym or not
        attention = 0.5 * (attention + attention.T)
        if self.is_randomwalk_laplacian:
            laplacian = calculate_randomwalk_normalized_laplacian(attention)
        else:
            laplacian = calculate_normalized_laplacian(attention).to(self.device)
        Lambda_, U = torch.symeig(laplacian, eigenvectors=True)
        return Lambda_, U, attention
//  [python]
def latent_correlation_layer(self, x):
        batch_size = x.shape[0]
        output, hidden = self.GRU(x)
        output = output[:, -1:, :]
        attention = self.self_graph_attention(output)
        attention = torch.mean(attention, dim=0)  # average of batch
        degree = torch.sum(attention, dim=1)
        # laplacian is sym or not
        attention = 0.5 * (attention + attention.T)
        if self.is_randomwalk_laplacian:
            laplacian = calculate_randomwalk_normalized_laplacian(attention)
        else:
            laplacian = calculate_normalized_laplacian(attention).to(self.device)
        Lambda_, U = torch.linalg.eigh(laplacian)
        return Lambda_, U, attention
//  [python]
def GDPPLoss(phiFake, phiReal, backward=True):
    r"""
    Implementation of the GDPP loss. Can be used with any kind of GAN
    architecture.

    Args:

        phiFake (tensor) : last feature layer of the discriminator on real data
        phiReal (tensor) : last feature layer of the discriminator on fake data
        backward (bool)  : should we perform the backward operation ?

    Returns:

        Loss's value. The backward operation in performed within this operator
    """
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.symeig(SB, eigenvectors=True)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        if abs(minV - maxV) < 1e-10:
            return eigVals
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phiFake)
    realEigVals, realEigVecs = compute_diversity(phiReal)

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()
//  [python]
def GDPPLoss(phiFake, phiReal, backward=True):
    r"""
    Implementation of the GDPP loss. Can be used with any kind of GAN
    architecture.

    Args:

        phiFake (tensor) : last feature layer of the discriminator on real data
        phiReal (tensor) : last feature layer of the discriminator on fake data
        backward (bool)  : should we perform the backward operation ?

    Returns:

        Loss's value. The backward operation in performed within this operator
    """
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.linalg.eigh(SB)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        if abs(minV - maxV) < 1e-10:
            return eigVals
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phiFake)
    realEigVals, realEigVecs = compute_diversity(phiReal)

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()
//  [python]
def reduceRank(A):
    r"""Reduces a symmetric factorization to a full rank symmetric factorization.

    Given a matrix :math:`U` of size :math`[n \times k]`, returns a matrix :math:`W`
    of size :math:`[n \times r(U)]`, where :math:`r(U)` is the rank of :math:`U`,
    such that :math:`WW`\top = UU^\top`.
    """
    
    (e, V) = torch.symeig(A.t() @ A, eigenvectors=True)
    
    eps = 10**-6
    
    Sigma = torch.sqrt(e)
    mask = 1-torch.isnan(Sigma)
    return A @ V[:, mask]
//  [python]
def reduceRank(A):
    r"""Reduces a symmetric factorization to a full rank symmetric factorization.

    Given a matrix :math:`U` of size :math`[n \times k]`, returns a matrix :math:`W`
    of size :math:`[n \times r(U)]`, where :math:`r(U)` is the rank of :math:`U`,
    such that :math:`WW`\top = UU^\top`.
    """
    
    (e, V) = torch.linalg.eigh(A.t() @ A)
    
    eps = 10**-6
    
    Sigma = torch.sqrt(e)
    mask = 1-torch.isnan(Sigma)
    return A @ V[:, mask]
//  [python]
def build(self):

        self.mean_x /= self.N_x
        self.covar_x /= self.N_x
        self.covar_x = self.covar_x - \
            self.mean_x.view(1, -1) * self.mean_x.view(-1, 1)

        # Variance around each dimension
        self.square_x /= self.N_x
        self.square_x = \
            torch.sqrt(torch.clamp(self.square_x - self.mean_x * self.mean_x,
                                   min=0))
        inv_square_x = 1 / (self.square_x + 1e-08)

        # Covariance matrix of the noramlized value of x
        covar_x_normalized = inv_square_x.view(-1, 1) * \
            self.covar_x * inv_square_x.view(1, -1)

        # Inverse Cholesky decomposition of the normalized matrix
        # l l.t = covar_x_normalized
        l_ = torch.cholesky(covar_x_normalized).inverse().detach()
        self.normalizer = l_.view(1, self.k, self.k).clone()

        # Now build the speed covariance matrix
        self.covar_speed /= self.N_speed
        self.covar_speed = inv_square_x.view(-1, 1) * \
            self.covar_speed * inv_square_x.view(1, -1)
        self.covar_speed = torch.mm(l_, torch.mm(self.covar_speed, l_.t()))

        e_vals, e_vects = torch.symeig(self.covar_speed, eigenvectors=True)
        k = e_vects.size(0)
        self.PCA_mul = e_vects.t().view(1, k, k).clone()
        self.PCA_values = e_vals.clone()
        self.building = False
        self.projection = self.PCA_mul.clone()
//  [python]
def build(self):

        self.mean_x /= self.N_x
        self.covar_x /= self.N_x
        self.covar_x = self.covar_x - \
            self.mean_x.view(1, -1) * self.mean_x.view(-1, 1)

        # Variance around each dimension
        self.square_x /= self.N_x
        self.square_x = \
            torch.sqrt(torch.clamp(self.square_x - self.mean_x * self.mean_x,
                                   min=0))
        inv_square_x = 1 / (self.square_x + 1e-08)

        # Covariance matrix of the noramlized value of x
        covar_x_normalized = inv_square_x.view(-1, 1) * \
            self.covar_x * inv_square_x.view(1, -1)

        # Inverse Cholesky decomposition of the normalized matrix
        # l l.t = covar_x_normalized
        l_ = torch.cholesky(covar_x_normalized).inverse().detach()
        self.normalizer = l_.view(1, self.k, self.k).clone()

        # Now build the speed covariance matrix
        self.covar_speed /= self.N_speed
        self.covar_speed = inv_square_x.view(-1, 1) * \
            self.covar_speed * inv_square_x.view(1, -1)
        self.covar_speed = torch.mm(l_, torch.mm(self.covar_speed, l_.t()))

        e_vals, e_vects = torch.linalg.eigh(self.covar_speed)
        k = e_vects.size(0)
        self.PCA_mul = e_vects.t().view(1, k, k).clone()
        self.PCA_values = e_vals.clone()
        self.building = False
        self.projection = self.PCA_mul.clone()
//  [python]
def forward(self,inp):
        """ (bs,n,d) (bs,n,c), (bs,n,n)"""

        xin,meanin,kin = inp
        e = torch.symeig(kin)[0]
        assert not (e < 1e-6).any(), (e.min(), kin.shape)
        if debug: assert not anynan(meanin)
        if debug: assert not anynan(kin)
        bs,n,c = meanin.shape
        kin += torch.eye(n,device=xin.device)*1e-6
        kdiag = torch.diagonal(kin,dim1=-2,dim2=-1) # (bs,n)
        stdin = kdiag.unsqueeze(-1).clamp(min=1e-8).sqrt() #(bs,n,1)
        # compute meanout:
        Phis = Phi(meanin/stdin)
        phis = phi(meanin/stdin)
        Eh = meanin * Phis + stdin * phis #(bs,n,c)
        # compute Kout
        K1 = torch.stack([kdiag[...,None].repeat((1,1,n)),kin],dim=-1) #(bs,n,n,2)
        K2 = torch.stack([kin.transpose(-1,-2),kdiag[...,None,:].repeat((1,n,1))],dim=-1)
        K12 = torch.stack([K1,K2],dim=-2) # (bs,n,n,2,2)
        #K12 += 10*torch.eye(n,device=xin.device)[:,:,None,None]*torch.eye(2,device=xin.device)[None,None]
        if debug: assert not anynan(K12)
        sqrtK12 = sqrt2x2(K12)
        #if debug: assert (sqrtK12@sqrtK12-K12).abs().mean()<1e-5
        isqrtK12 = inverse2x2(sqrtK12) #(bs,n,n,2,2)

        if debug: assert not anynan(sqrtK12)
        if debug: assert not anynan(isqrtK12)
        u1 = isqrtK12[...,0,0,None]*meanin[:,:,None,:]+isqrtK12[...,0,1,None]*meanin[:,None,:,:] #(bs,n,n,c)
        u2 = isqrtK12[...,1,0,None]*meanin[:,:,None,:]+isqrtK12[...,1,1,None]*meanin[:,None,:,:]
        Phi1 = Phi(u1)
        Phi2 = Phi(u2)
        Psi = Phi1*Phi2
        psi1 = phi(u1)*Phi2
        psi2 = Phi1*phi(u2)
        EhhT_Psi = (meanin[:,:,None,:]*meanin[:,None,:,:]+kin[...,None])*Psi
        # mu2S1 + mu1S2
        EhhT_psi1 = (meanin[:,None,:,:]*sqrtK12[...,0,0,None]+meanin[:,:,None,:]*sqrtK12[...,1,0,None])*psi1
        EhhT_psi2 = (meanin[:,None,:,:]*sqrtK12[...,0,1,None]+meanin[:,:,None,:]*sqrtK12[...,1,1,None])*psi2
        SHPhiS11 = sqrtK12[...,0,0,None]*sqrtK12[...,0,1,None]*phip(u1)*Phi(u2)
        SHPhiS12 = sqrtK12[..., 0, 1,None] * sqrtK12[..., 0, 1,None] * phi(u1) * phi(u2)
        SHPhiS22 = sqrtK12[..., 0, 1,None] * sqrtK12[..., 1, 1,None] * Phi(u1) * phip(u2)
        SHPhiS = SHPhiS11+2*SHPhiS12+SHPhiS22
        EhhT = EhhT_Psi + (EhhT_psi1+EhhT_psi2) + SHPhiS
        covar = (EhhT - Eh[:,:,None,:]*Eh[:,None,:,:]).mean(-1) #(bs,n,n)
        diag_covar = torch.diag_embed(((meanin**2 + stdin**2)*Phis+meanin*stdin*phis - Eh**2).mean(-1).clamp(min=1e-6)) #(bs,n,n)
        full_covar = torch.where(diag_covar>1e-7,diag_covar,covar)
        e = torch.symeig(full_covar)[0]
        assert not (e<1e-6).any(), (e.min(),full_covar.shape)
        return xin,Eh,full_covar+torch.eye(n,device=xin.device)*1e-6
//  [python]
def forward(self,inp):
        """ (bs,n,d) (bs,n,c), (bs,n,n)"""

        xin,meanin,kin = inp
        e = torch.linalg.eigh(kin)[0]
        assert not (e < 1e-6).any(), (e.min(), kin.shape)
        if debug: assert not anynan(meanin)
        if debug: assert not anynan(kin)
        bs,n,c = meanin.shape
        kin += torch.eye(n,device=xin.device)*1e-6
        kdiag = torch.diagonal(kin,dim1=-2,dim2=-1) # (bs,n)
        stdin = kdiag.unsqueeze(-1).clamp(min=1e-8).sqrt() #(bs,n,1)
        # compute meanout:
        Phis = Phi(meanin/stdin)
        phis = phi(meanin/stdin)
        Eh = meanin * Phis + stdin * phis #(bs,n,c)
        # compute Kout
        K1 = torch.stack([kdiag[...,None].repeat((1,1,n)),kin],dim=-1) #(bs,n,n,2)
        K2 = torch.stack([kin.transpose(-1,-2),kdiag[...,None,:].repeat((1,n,1))],dim=-1)
        K12 = torch.stack([K1,K2],dim=-2) # (bs,n,n,2,2)
        #K12 += 10*torch.eye(n,device=xin.device)[:,:,None,None]*torch.eye(2,device=xin.device)[None,None]
        if debug: assert not anynan(K12)
        sqrtK12 = sqrt2x2(K12)
        #if debug: assert (sqrtK12@sqrtK12-K12).abs().mean()<1e-5
        isqrtK12 = inverse2x2(sqrtK12) #(bs,n,n,2,2)

        if debug: assert not anynan(sqrtK12)
        if debug: assert not anynan(isqrtK12)
        u1 = isqrtK12[...,0,0,None]*meanin[:,:,None,:]+isqrtK12[...,0,1,None]*meanin[:,None,:,:] #(bs,n,n,c)
        u2 = isqrtK12[...,1,0,None]*meanin[:,:,None,:]+isqrtK12[...,1,1,None]*meanin[:,None,:,:]
        Phi1 = Phi(u1)
        Phi2 = Phi(u2)
        Psi = Phi1*Phi2
        psi1 = phi(u1)*Phi2
        psi2 = Phi1*phi(u2)
        EhhT_Psi = (meanin[:,:,None,:]*meanin[:,None,:,:]+kin[...,None])*Psi
        # mu2S1 + mu1S2
        EhhT_psi1 = (meanin[:,None,:,:]*sqrtK12[...,0,0,None]+meanin[:,:,None,:]*sqrtK12[...,1,0,None])*psi1
        EhhT_psi2 = (meanin[:,None,:,:]*sqrtK12[...,0,1,None]+meanin[:,:,None,:]*sqrtK12[...,1,1,None])*psi2
        SHPhiS11 = sqrtK12[...,0,0,None]*sqrtK12[...,0,1,None]*phip(u1)*Phi(u2)
        SHPhiS12 = sqrtK12[..., 0, 1,None] * sqrtK12[..., 0, 1,None] * phi(u1) * phi(u2)
        SHPhiS22 = sqrtK12[..., 0, 1,None] * sqrtK12[..., 1, 1,None] * Phi(u1) * phip(u2)
        SHPhiS = SHPhiS11+2*SHPhiS12+SHPhiS22
        EhhT = EhhT_Psi + (EhhT_psi1+EhhT_psi2) + SHPhiS
        covar = (EhhT - Eh[:,:,None,:]*Eh[:,None,:,:]).mean(-1) #(bs,n,n)
        diag_covar = torch.diag_embed(((meanin**2 + stdin**2)*Phis+meanin*stdin*phis - Eh**2).mean(-1).clamp(min=1e-6)) #(bs,n,n)
        full_covar = torch.where(diag_covar>1e-7,diag_covar,covar)
        e = torch.linalg.eigh(full_covar)[0]
        assert not (e<1e-6).any(), (e.min(),full_covar.shape)
        return xin,Eh,full_covar+torch.eye(n,device=xin.device)*1e-6
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    #print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  #(B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  #(B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  #(B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask>1] = 1
        
        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.symeig(L, eigenvectors=True)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    #print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  #(B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  #(B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  #(B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask>1] = 1
        
        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.linalg.eigh(L)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def overlap(self, psi):
        '''
        overlap of the ground state with the given state (assuming normalized)
        '''
        w, v = torch.symeig(self._buildH(), eigenvectors=True)
        return (v[:, 0] * psi).sum().abs()
//  [python]
def overlap(self, psi):
        '''
        overlap of the ground state with the given state (assuming normalized)
        '''
        w, v = torch.linalg.eigh(self._buildH())
        return (v[:, 0] * psi).sum().abs()
//  [python]
def test_symeig(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_hermitian_matrix(*dims, dtype=dtype, device=device)
            if dtype.is_complex:
                real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
            else:
                real_dtype = dtype
            oute = torch.empty(dims[1:] + dims[:1], dtype=real_dtype, device=device)
            outv = torch.empty(dims[1:] + dims[:1] * 2, dtype=dtype, device=device)
            torch.symeig(x, eigenvectors=eigenvectors, upper=upper, out=(oute, outv))

            if eigenvectors:
                outv_ = outv.cpu().numpy()
                x_recon = np.matmul(np.matmul(outv_, torch.diag_embed(oute.to(dtype)).cpu().numpy()),
                                    outv_.swapaxes(-2, -1).conj())
                self.assertEqual(x, x_recon, atol=1e-8, rtol=0, msg='Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.symeig(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, oute, msg='Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), outv, msg='Eigenvector matrix not empty')

            rese, resv = x.symeig(eigenvectors=eigenvectors, upper=upper)
            self.assertEqual(rese, oute, msg="outputs of symeig and symeig with out don't match")
            self.assertEqual(resv, outv, msg="outputs of symeig and symeig with out don't match")

            # test non-contiguous
            x = random_hermitian_matrix(*dims, dtype=dtype, device=device)
            n_dim = len(dims) + 1
            # Reverse the batch dimensions and the matrix dimensions and then concat them
            x = x.permute(tuple(range(n_dim - 3, -1, -1)) + (n_dim - 1, n_dim - 2))
            assert not x.is_contiguous(), "x is intentionally non-contiguous"
            rese, resv = torch.symeig(x, eigenvectors=eigenvectors, upper=upper)
            if eigenvectors:
                resv_ = resv.cpu().numpy()
                x_recon = np.matmul(np.matmul(resv_, torch.diag_embed(rese.to(dtype)).cpu().numpy()),
                                    resv_.swapaxes(-2, -1).conj())
                self.assertEqual(x, x_recon, atol=1e-8, rtol=0, msg='Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.symeig(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, rese, msg='Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), resv, msg='Eigenvector matrix not empty')

        batch_dims_set = [(), (3,), (3, 5), (5, 3, 5)]
        for batch_dims, eigenvectors, upper in itertools.product(batch_dims_set, (True, False), (True, False)):
            run_test((5,) + batch_dims, eigenvectors, upper)
//  [python]
def test_symeig(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_hermitian_matrix(*dims, dtype=dtype, device=device)
            if dtype.is_complex:
                real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
            else:
                real_dtype = dtype
            oute = torch.empty(dims[1:] + dims[:1], dtype=real_dtype, device=device)
            outv = torch.empty(dims[1:] + dims[:1] * 2, dtype=dtype, device=device)
            torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper, out=(oute, outv))

            if eigenvectors:
                outv_ = outv.cpu().numpy()
                x_recon = np.matmul(np.matmul(outv_, torch.diag_embed(oute.to(dtype)).cpu().numpy()),
                                    outv_.swapaxes(-2, -1).conj())
                self.assertEqual(x, x_recon, atol=1e-8, rtol=0, msg='Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.linalg.eigh(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, oute, msg='Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), outv, msg='Eigenvector matrix not empty')

            rese, resv = torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper)
            self.assertEqual(rese, oute, msg="outputs of symeig and symeig with out don't match")
            self.assertEqual(resv, outv, msg="outputs of symeig and symeig with out don't match")

            # test non-contiguous
            x = random_hermitian_matrix(*dims, dtype=dtype, device=device)
            n_dim = len(dims) + 1
            # Reverse the batch dimensions and the matrix dimensions and then concat them
            x = x.permute(tuple(range(n_dim - 3, -1, -1)) + (n_dim - 1, n_dim - 2))
            assert not x.is_contiguous(), "x is intentionally non-contiguous"
            rese, resv = torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper)
            if eigenvectors:
                resv_ = resv.cpu().numpy()
                x_recon = np.matmul(np.matmul(resv_, torch.diag_embed(rese.to(dtype)).cpu().numpy()),
                                    resv_.swapaxes(-2, -1).conj())
                self.assertEqual(x, x_recon, atol=1e-8, rtol=0, msg='Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.linalg.eigh(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, rese, msg='Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), resv, msg='Eigenvector matrix not empty')

        batch_dims_set = [(), (3,), (3, 5), (5, 3, 5)]
        for batch_dims, eigenvectors, upper in itertools.product(batch_dims_set, (True, False), (True, False)):
            run_test((5,) + batch_dims, eigenvectors, upper)
//  [python]
def kway_normcuts(F, K=2, sigma=1.0):
    # Build similarity matrix W, use cosine similarity
    W = embeddings_to_cosine_similarity(F, sigma=sigma)

    # Build defree matrix
    degree = torch.sum(W, dim=0)

    # Construct normalized Laplacian matrix L
    D_pow = torch.diag(degree.pow(-0.5))
    L = torch.matmul(torch.matmul(D_pow, torch.diag(degree)-W), D_pow)

    # Get eigenvectors with torch.symeig()
    _, eigenvector = torch.symeig(L, eigenvectors=True)

    # Normalize eigenvector along each row
    eigvec_norm = torch.matmul(torch.diag(degree.pow(-0.5)), eigenvector)
    eigvec_norm = eigvec_norm/eigvec_norm[0][0]
    k_eigvec = eigvec_norm[:,:K]

    return k_eigvec
//  [python]
def kway_normcuts(F, K=2, sigma=1.0):
    # Build similarity matrix W, use cosine similarity
    W = embeddings_to_cosine_similarity(F, sigma=sigma)

    # Build defree matrix
    degree = torch.sum(W, dim=0)

    # Construct normalized Laplacian matrix L
    D_pow = torch.diag(degree.pow(-0.5))
    L = torch.matmul(torch.matmul(D_pow, torch.diag(degree)-W), D_pow)

    # Get eigenvectors with torch.linalg.eigh
    _, eigenvector = torch.linalg.eigh(L)

    # Normalize eigenvector along each row
    eigvec_norm = torch.matmul(torch.diag(degree.pow(-0.5)), eigenvector)
    eigvec_norm = eigvec_norm/eigvec_norm[0][0]
    k_eigvec = eigvec_norm[:,:K]

    return k_eigvec
//  [python]
def _update_inv(self):
        assert self.steps > 0, 'At least one step before update inverse!'
        eps = 1e-10
        for idx, m in enumerate(self.modules):
            m_aa, m_gg = self.m_aa[m], self.m_gg[m]
            self.d_a[m], Q_a = torch.symeig(m_aa, eigenvectors=True)
            self.d_g[m], Q_g = torch.symeig(m_gg, eigenvectors=True)

            self.d_a[m].mul_((self.d_a[m] > eps).float())
            self.d_g[m].mul_((self.d_g[m] > eps).float())

            # == write summary ==
            name = m.__class__.__name__
            eigs = (self.d_g[m].view(-1, 1) @ self.d_a[m].view(1, -1)).view(-1).cpu().data.numpy()
            self.writer.add_histogram('eigen/%s_%d' % (name, idx), eigs, self.iter)

            if self.Q_a.get(m, None) is None:
                # print('(%d)Q_a %s is None.' % (idx, m))
                self.Q_a[m] = [Q_a]  # absorb the eigen basis
            else:
                # self.Q_a[m] = [Q_a, self.Q_a[m]]
                prev_Q_a, prev_Q_g = get_rotation_layer_weights(self.model, m)
                prev_Q_a = prev_Q_a.view(prev_Q_a.size(0), prev_Q_a.size(1)).transpose(1, 0)
                prev_Q_g = prev_Q_g.view(prev_Q_g.size(0), prev_Q_g.size(1))
                self.Q_a[m] = [Q_a, prev_Q_a]

            if self.Q_g.get(m, None) is None:
                self.Q_g[m] = [Q_g]
            else:
                self.Q_g[m] = [Q_g, prev_Q_g]
        self._inversed = True
        self.iter += 1
//  [python]
def _update_inv(self):
        assert self.steps > 0, 'At least one step before update inverse!'
        eps = 1e-10
        for idx, m in enumerate(self.modules):
            m_aa, m_gg = self.m_aa[m], self.m_gg[m]
            self.d_a[m], Q_a = torch.linalg.eigh(m_aa)
            self.d_g[m], Q_g = torch.linalg.eigh(m_gg)

            self.d_a[m].mul_((self.d_a[m] > eps).float())
            self.d_g[m].mul_((self.d_g[m] > eps).float())

            # == write summary ==
            name = m.__class__.__name__
            eigs = (self.d_g[m].view(-1, 1) @ self.d_a[m].view(1, -1)).view(-1).cpu().data.numpy()
            self.writer.add_histogram('eigen/%s_%d' % (name, idx), eigs, self.iter)

            if self.Q_a.get(m, None) is None:
                # print('(%d)Q_a %s is None.' % (idx, m))
                self.Q_a[m] = [Q_a]  # absorb the eigen basis
            else:
                # self.Q_a[m] = [Q_a, self.Q_a[m]]
                prev_Q_a, prev_Q_g = get_rotation_layer_weights(self.model, m)
                prev_Q_a = prev_Q_a.view(prev_Q_a.size(0), prev_Q_a.size(1)).transpose(1, 0)
                prev_Q_g = prev_Q_g.view(prev_Q_g.size(0), prev_Q_g.size(1))
                self.Q_a[m] = [Q_a, prev_Q_a]

            if self.Q_g.get(m, None) is None:
                self.Q_g[m] = [Q_g]
            else:
                self.Q_g[m] = [Q_g, prev_Q_g]
        self._inversed = True
        self.iter += 1
//  [python]
def compute_principal_curvatures(beta):
    """
    given the jet coefficients, compute the principal curvatures and principal directions:
    the eigenvalues and eigenvectors of the weingarten matrix
    :param beta: batch of Jet coefficients vector
    :return: k1, k2, dir1, dir2: batch of principal curvatures and principal directions
    """
    with torch.no_grad():
        if beta.shape[1] < 5:
            raise ValueError("Can't compute curvatures for jet with order less than 2")
        else:
            b1_2 = torch.pow(beta[:, 0], 2)
            b2_2 = torch.pow(beta[:, 1], 2)
            #first fundemental form
            E = (1 + b1_2).view(-1, 1, 1)
            G = (1 + b2_2).view(-1, 1, 1)
            F = (beta[:, 1] * beta[:, 0]).view(-1, 1, 1)
            I = torch.cat([torch.cat([E, F], dim=2), torch.cat([F, G], dim=2)], dim=1)
            # second fundemental form
            norm_N0 = torch.sqrt(b1_2 + b2_2 + 1)
            e = (2*beta[:, 2] / norm_N0).view(-1, 1, 1)
            f = (beta[:, 4] / norm_N0).view(-1, 1, 1)
            g = (2*beta[:, 3] / norm_N0).view(-1, 1, 1)
            II = torch.cat([torch.cat([e, f], dim=2), torch.cat([f, g], dim=2)], dim=1)

            M_weingarten = -torch.bmm(torch.inverse(I), II)

            curvatures, dirs = torch.symeig(M_weingarten, eigenvectors=True) #faster
            dirs = torch.cat([dirs, torch.zeros(dirs.shape[0], 2, 1, device=dirs.device)], dim=2) # pad zero in the normal direction

    return curvatures, dirs
//  [python]
def compute_principal_curvatures(beta):
    """
    given the jet coefficients, compute the principal curvatures and principal directions:
    the eigenvalues and eigenvectors of the weingarten matrix
    :param beta: batch of Jet coefficients vector
    :return: k1, k2, dir1, dir2: batch of principal curvatures and principal directions
    """
    with torch.no_grad():
        if beta.shape[1] < 5:
            raise ValueError("Can't compute curvatures for jet with order less than 2")
        else:
            b1_2 = torch.pow(beta[:, 0], 2)
            b2_2 = torch.pow(beta[:, 1], 2)
            #first fundemental form
            E = (1 + b1_2).view(-1, 1, 1)
            G = (1 + b2_2).view(-1, 1, 1)
            F = (beta[:, 1] * beta[:, 0]).view(-1, 1, 1)
            I = torch.cat([torch.cat([E, F], dim=2), torch.cat([F, G], dim=2)], dim=1)
            # second fundemental form
            norm_N0 = torch.sqrt(b1_2 + b2_2 + 1)
            e = (2*beta[:, 2] / norm_N0).view(-1, 1, 1)
            f = (beta[:, 4] / norm_N0).view(-1, 1, 1)
            g = (2*beta[:, 3] / norm_N0).view(-1, 1, 1)
            II = torch.cat([torch.cat([e, f], dim=2), torch.cat([f, g], dim=2)], dim=1)

            M_weingarten = -torch.bmm(torch.inverse(I), II)

            curvatures, dirs = torch.linalg.eigh(M_weingarten)
            dirs = torch.cat([dirs, torch.zeros(dirs.shape[0], 2, 1, device=dirs.device)], dim=2) # pad zero in the normal direction

    return curvatures, dirs
//  [python]
parametrize.register_parametrization(layer_spd, "weight", Symmetric())
parametrize.register_parametrization(layer_spd, "weight", MatrixExponential())
X = layer_spd.weight
print(torch.dist(X, X.T))                        # X is symmetric
print((torch.symeig(X).eigenvalues > 0.).all())
//  [python]
parametrize.register_parametrization(layer_spd, "weight", Symmetric())
parametrize.register_parametrization(layer_spd, "weight", MatrixExponential())
X = layer_spd.weight
print(torch.dist(X, X.T))                        # X is symmetric
print((torch.linalg.eigh(X).eigenvalues > 0.).all())
//  [python]
def chen_estimate(im,pch_size=8):
    im=torch.squeeze(im)
    
    #grayscale
    im=im.unsqueeze(0)
    pch=im2patch(im,pch_size,3)
    num_pch=pch.size()[3]
    pch=pch.view((-1,num_pch))
    d=pch.size()[0]
    mu=torch.mean(pch,dim=1,keepdim=True)
    
    X=pch-mu
    sigma_X=torch.matmul(X,torch.t(X))/num_pch
    sig_value,_=torch.symeig(sigma_X,eigenvectors=True)
    sig_value=sig_value.sort().values
    
    
    start=time.time()
    # tensor operation for substituting iterative step.
    # These operation make  parallel computing possiblie which is more efficient

    triangle=torch.ones((d,d))
    triangle= torch.tril(triangle).cuda()
    sig_matrix= torch.matmul( triangle, torch.diag(sig_value)) 
    
    # calculate whole threshold value at a single time
    num_vec= torch.arange(d)+1
    num_vec=num_vec.to(dtype=torch.float32).cuda()
    sum_arr= torch.sum(sig_matrix,dim=1)
    tau_arr=sum_arr/num_vec
    
    tau_mat= torch.matmul(torch.diag(tau_arr),triangle)
    
    # find median value with masking scheme: 
    big_bool= torch.sum(sig_matrix>tau_mat,axis=1)
    small_bool= torch.sum(sig_matrix<tau_mat,axis=1)
    mask=(big_bool==small_bool).to(dtype=torch.float32).cuda()
    tau_chen=torch.max(mask*tau_arr)
      
# Previous implementation       
#    for ii in range(-1, -d-1, -1):
#        tau = torch.mean(sig_value[:ii])
#        if torch.sum(sig_value[:ii]>tau) == torch.sum(sig_value[:ii] < tau):
             #  return torch.sqrt(tau)
#    print('old: ', torch.sqrt(tau))

    return torch.sqrt(tau_chen)
//  [python]
def chen_estimate(im,pch_size=8):
    im=torch.squeeze(im)
    
    #grayscale
    im=im.unsqueeze(0)
    pch=im2patch(im,pch_size,3)
    num_pch=pch.size()[3]
    pch=pch.view((-1,num_pch))
    d=pch.size()[0]
    mu=torch.mean(pch,dim=1,keepdim=True)
    
    X=pch-mu
    sigma_X=torch.matmul(X,torch.t(X))/num_pch
    sig_value,_=torch.linalg.eigh(sigma_X)
    sig_value=sig_value.sort().values
    
    
    start=time.time()
    # tensor operation for substituting iterative step.
    # These operation make  parallel computing possiblie which is more efficient

    triangle=torch.ones((d,d))
    triangle= torch.tril(triangle).cuda()
    sig_matrix= torch.matmul( triangle, torch.diag(sig_value)) 
    
    # calculate whole threshold value at a single time
    num_vec= torch.arange(d)+1
    num_vec=num_vec.to(dtype=torch.float32).cuda()
    sum_arr= torch.sum(sig_matrix,dim=1)
    tau_arr=sum_arr/num_vec
    
    tau_mat= torch.matmul(torch.diag(tau_arr),triangle)
    
    # find median value with masking scheme: 
    big_bool= torch.sum(sig_matrix>tau_mat,axis=1)
    small_bool= torch.sum(sig_matrix<tau_mat,axis=1)
    mask=(big_bool==small_bool).to(dtype=torch.float32).cuda()
    tau_chen=torch.max(mask*tau_arr)
      
# Previous implementation       
#    for ii in range(-1, -d-1, -1):
#        tau = torch.mean(sig_value[:ii])
#        if torch.sum(sig_value[:ii]>tau) == torch.sum(sig_value[:ii] < tau):
             #  return torch.sqrt(tau)
#    print('old: ', torch.sqrt(tau))

    return torch.sqrt(tau_chen)
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.symeig(matrix, eigenvectors=True)
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.linalg.eigh(matrix)
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def compute_eigendecomposition(self, impl='symeig'):
        # TODO: test
        if impl == 'symeig':
            self.evals, self.evecs = torch.symeig(self.data, eigenvectors=True)
        elif impl == 'svd':
            _, self.evals, self.evecs = torch.svd(self.data, some=False)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl='symeig'):
        if impl == 'symeig':
            self.evals, self.evecs = torch.linalg.eigh(self.data)
        elif impl == 'svd':
            _, self.evals, self.evecs = torch.svd(self.data, some=False)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl='symeig'):
        self.evals = dict()
        self.evecs = dict()
        if impl == 'symeig':
            for layer_id in self.generator.layer_collection.layers.keys():
                a, g = self.data[layer_id]
                evals_a, evecs_a = torch.symeig(a, eigenvectors=True)
                evals_g, evecs_g = torch.symeig(g, eigenvectors=True)
                self.evals[layer_id] = (evals_a, evals_g)
                self.evecs[layer_id] = (evecs_a, evecs_g)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl='symeig'):
        self.evals = dict()
        self.evecs = dict()
        if impl == 'symeig':
            for layer_id in self.generator.layer_collection.layers.keys():
                a, g = self.data[layer_id]
                evals_a, evecs_a = torch.linalg.eigh(a)
                evals_g, evecs_g = torch.linalg.eigh(g)
                self.evals[layer_id] = (evals_a, evals_g)
                self.evecs[layer_id] = (evecs_a, evecs_g)
        else:
            raise NotImplementedError
//  [python]
def symeigLanczos(A, k, device=torch.device("cpu"), extreme="both", *, 
                    sparse=False, dim=None):
    """
        This function computes the extreme(minimum or maximum, or both) eigenvalues
    and corresponding eigenvectors of a real symmetric matrix A based on Lanczos algorithm.

    Input: `A` is the real symmetric matrix to be diagonalized.
           `k` is the number of Lanczos vectors requested.
           `sparse` indicates whether a bare linear function representation of the
                matrix A is adopted. When it is `true`, the integer parameter `dim`
                must be supplied giving the actual dimension of square matrix A.
           `extreme` labels which of the two extreme eigenvalues and corresponding eigenvectors are returned.
            "both" -> both min and max.     --Output--> (eigval_min, eigvector_min, eigval_max, eigvector_max)
            "min" -> min.                   --Output--> (eigval_min, eigvector_min)
            "max" -> max.                   --Output--> (eigval_max, eigvector_max)
    Output: As shown in "Input" section above. Note all the elements of
            returned tuples are torch Tensors, including the eigenvalues.
    """
    Qk, T = Lanczos(A, k, device=device, sparse=sparse, dim=dim)
    eigvalsQ, eigvectorsQ = torch.symeig(T, eigenvectors=True)
    eigvectorsQ = torch.matmul(Qk, eigvectorsQ)
    if extreme == "both":
        return eigvalsQ[0], eigvectorsQ[:, 0], eigvalsQ[-1], eigvectorsQ[:, -1]
    elif extreme == "min":
        return eigvalsQ[0], eigvectorsQ[:, 0]
    elif extreme == "max":
        return eigvalsQ[-1], eigvectorsQ[:, -1]
//  [python]
def symeigLanczos(A, k, device=torch.device("cpu"), extreme="both", *, 
                    sparse=False, dim=None):
    """
        This function computes the extreme(minimum or maximum, or both) eigenvalues
    and corresponding eigenvectors of a real symmetric matrix A based on Lanczos algorithm.

    Input: `A` is the real symmetric matrix to be diagonalized.
           `k` is the number of Lanczos vectors requested.
           `sparse` indicates whether a bare linear function representation of the
                matrix A is adopted. When it is `true`, the integer parameter `dim`
                must be supplied giving the actual dimension of square matrix A.
           `extreme` labels which of the two extreme eigenvalues and corresponding eigenvectors are returned.
            "both" -> both min and max.     --Output--> (eigval_min, eigvector_min, eigval_max, eigvector_max)
            "min" -> min.                   --Output--> (eigval_min, eigvector_min)
            "max" -> max.                   --Output--> (eigval_max, eigvector_max)
    Output: As shown in "Input" section above. Note all the elements of
            returned tuples are torch Tensors, including the eigenvalues.
    """
    Qk, T = Lanczos(A, k, device=device, sparse=sparse, dim=dim)
    eigvalsQ, eigvectorsQ = torch.linalg.eigh(T)
    eigvectorsQ = torch.matmul(Qk, eigvectorsQ)
    if extreme == "both":
        return eigvalsQ[0], eigvectorsQ[:, 0], eigvalsQ[-1], eigvectorsQ[:, -1]
    elif extreme == "min":
        return eigvalsQ[0], eigvectorsQ[:, 0]
    elif extreme == "max":
        return eigvalsQ[-1], eigvectorsQ[:, -1]