jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def evd(A):
    """Matrix square root and inverse square root for SPD matrices"""
    ei, ev = torch.symeig(A, eigenvectors=True)
    ei = torch.diag(torch.sqrt(ei))
    ei_inv = torch.diag(1. / torch.sqrt(ei))
    Asqrt = ev @ ei @ ev.t()
    Asqrt_inv = ev @ ei_inv @ ev.t()
    return Asqrt, Asqrt_inv
//  [python]
def evd(A):
    """Matrix square root and inverse square root for SPD matrices"""
    ei, ev = torch.linalg.eigh(A)
    ei = torch.diag(torch.sqrt(ei))
    ei_inv = torch.diag(1. / torch.sqrt(ei))
    Asqrt = ev @ ei @ ev.t()
    Asqrt_inv = ev @ ei_inv @ ev.t()
    return Asqrt, Asqrt_inv
//  [python]
def pcawhitenlearn_torch(X, min_eig=-1):

    N = X.shape[1]

    # Learning PCA w/o annotations
    m = X.mean(axis=1, keepdims=True)
    Xc = X - m
    Xcov = torch.mm(Xc, Xc.T)
    Xcov = (Xcov + Xcov.T) / (2*N)
    eigval, eigvec = torch.symeig(Xcov, eigenvectors=True)
    order = torch.argsort(eigval, descending=True)
    eigval = eigval[order]
    eigvec = eigvec[:, order]

    beta = eigval[min_eig]
    P = torch.mm(torch.inverse(torch.sqrt(torch.diag((1-beta)*eigval + beta))), eigvec.T)
    return m, P
//  [python]
def pcawhitenlearn_torch(X, min_eig=-1):

    N = X.shape[1]

    # Learning PCA w/o annotations
    m = X.mean(axis=1, keepdims=True)
    Xc = X - m
    Xcov = torch.mm(Xc, Xc.T)
    Xcov = (Xcov + Xcov.T) / (2*N)
    eigval, eigvec = torch.linalg.eigh(Xcov)
    order = torch.argsort(eigval, descending=True)
    eigval = eigval[order]
    eigvec = eigvec[:, order]

    beta = eigval[min_eig]
    P = torch.mm(torch.inverse(torch.sqrt(torch.diag((1-beta)*eigval + beta))), eigvec.T)
    return m, P
//  [python]
def get_point_cluster(self,point):
        #Create Point Cloud
        num_point=point.shape[0]
        point_cloud=o3d.geometry.PointCloud()
        point_cloud.points=o3d.utility.Vector3dVector(point)

        #Downsample PointCloud into Voxel
        voxel=o3d.geometry.voxel_down_sample(point_cloud,0.1)
        voxel_coor=np.array(voxel.points)
        
        # Get Each Voxel's coordinate covairance matrix
        neighbor_size=20
        num_voxel=voxel_coor.shape[0]
        dist_matrix=spa.distance_matrix(voxel_coor,voxel_coor)
        index=np.argsort(dist_matrix,1)
        index=index[:,:neighbor_size]
        index=index.reshape(-1)
        normalized_coor=voxel_coor[index].reshape(num_voxel,neighbor_size,-1)
        normalized_coor=normalized_coor-np.mean(normalized_coor,1,keepdims=True)
        neighbor_coor=torch.FloatTensor(normalized_coor)
        coor_mat=torch.bmm(neighbor_coor.permute(0,2,1),neighbor_coor)
        
        ## get each feature cluster's coordinate ###
        e,v = torch.symeig(coor_mat, eigenvectors=True)
        labels = self.get_cluster_result(e,v)
        # coor_list=[]
        coor_list=np.zeros((0,50,3))
        for l in np.unique(labels):
            l_index=np.where(labels==l)[0]
            picked_ind=np.random.permutation(l_index)[:50]
            while len(picked_ind)<50:
                make_up=50-len(picked_ind)
                make_up_index=np.random.permutation(l_index)[:make_up]
                picked_ind=np.append(picked_ind,make_up_index)
            
            vo_co=np.expand_dims(voxel_coor[picked_ind],0)
            
            coor_list=np.append(coor_list,vo_co,0)
            
            # coor_list.append(voxel_coor[picked_ind])
        # coor_list=np.array(coor_list)
        return coor_list
//  [python]
def get_point_cluster(self,point):
        #Create Point Cloud
        num_point=point.shape[0]
        point_cloud=o3d.geometry.PointCloud()
        point_cloud.points=o3d.utility.Vector3dVector(point)

        #Downsample PointCloud into Voxel
        voxel=o3d.geometry.voxel_down_sample(point_cloud,0.1)
        voxel_coor=np.array(voxel.points)
        
        # Get Each Voxel's coordinate covairance matrix
        neighbor_size=20
        num_voxel=voxel_coor.shape[0]
        dist_matrix=spa.distance_matrix(voxel_coor,voxel_coor)
        index=np.argsort(dist_matrix,1)
        index=index[:,:neighbor_size]
        index=index.reshape(-1)
        normalized_coor=voxel_coor[index].reshape(num_voxel,neighbor_size,-1)
        normalized_coor=normalized_coor-np.mean(normalized_coor,1,keepdims=True)
        neighbor_coor=torch.FloatTensor(normalized_coor)
        coor_mat=torch.bmm(neighbor_coor.permute(0,2,1),neighbor_coor)
        
        ## get each feature cluster's coordinate ###
        e,v = torch.linalg.eigh(coor_mat)
        labels = self.get_cluster_result(e,v)
        # coor_list=[]
        coor_list=np.zeros((0,50,3))
        for l in np.unique(labels):
            l_index=np.where(labels==l)[0]
            picked_ind=np.random.permutation(l_index)[:50]
            while len(picked_ind)<50:
                make_up=50-len(picked_ind)
                make_up_index=np.random.permutation(l_index)[:make_up]
                picked_ind=np.append(picked_ind,make_up_index)
            
            vo_co=np.expand_dims(voxel_coor[picked_ind],0)
            
            coor_list=np.append(coor_list,vo_co,0)
            
            # coor_list.append(voxel_coor[picked_ind])
        # coor_list=np.array(coor_list)
        return coor_list
//  [python]
def whitenlearn_torch(X, qidxs, pidxs):

    # Learning Lw w annotations
    # m = X.mean(axis=1, keepdims=True)
    m = X[:, qidxs].mean(axis=1, keepdims=True)
    df = X[:, qidxs] - X[:, pidxs]
    S = torch.mm(df, df.T) / df.shape[1]
    P = torch.inverse(cholesky_torch(S))
    df = torch.mm(P, X-m)
    D = torch.mm(df, df.T)
    eigval, eigvec = torch.symeig(D, eigenvectors=True)
    order = torch.argsort(eigval, descending=True)
    eigval = eigval[order]
    eigvec = eigvec[:, order]

    P = torch.mm(eigvec.T, P)

    return m, P
//  [python]
def whitenlearn_torch(X, qidxs, pidxs):

    # Learning Lw w annotations
    # m = X.mean(axis=1, keepdims=True)
    m = X[:, qidxs].mean(axis=1, keepdims=True)
    df = X[:, qidxs] - X[:, pidxs]
    S = torch.mm(df, df.T) / df.shape[1]
    P = torch.inverse(cholesky_torch(S))
    df = torch.mm(P, X-m)
    D = torch.mm(df, df.T)
    eigval, eigvec = torch.linalg.eigh(D)
    order = torch.argsort(eigval, descending=True)
    eigval = eigval[order]
    eigvec = eigvec[:, order]

    P = torch.mm(eigvec.T, P)

    return m, P
//  [python]
def init_xyz_edm(mu,sigma,batch=1,steps=50):
    '''generate initial Ca trace from mu,sigma of 
    predicted log-normal distributions'''
    device = mu.device
    L = mu.shape[-1]
    m = mu.detach()
    s = sigma.detach()
    
    # approximate initial distances by the mode
    # of lognormal distributions
    D = torch.exp(m-s*s)
    sampler = LogNormal(mu[0],sigma[0])
    Ds = sampler.sample((batch,))
    Ds[0] = D[0]
    
    # convert distance matrix to a Gramm matrix
    J = torch.eye(L,device=device)-torch.ones((L,L),device=device)/L
    G = -0.5*J@Ds**2@J
    
    # use three leading eigenvectors to approximate
    # initial coordinates
    e,v = torch.symeig(G,eigenvectors=True)
    ca = v[:,:,-3:]*e[:,-3:].unsqueeze(1)**0.5
    
    def nll(xyz):
        dist = torch.cdist(xyz,xyz)+1e-6*torch.eye(L,device=device)
        nll = (0.5*((torch.log(dist)-m)/s + torch.log(dist*s*(2*np.pi)**0.5))**2).mean()
        return nll
    
    nll0 = nll(ca)
    nll1 = nll0
    
    # fine-tuning by minimization of log-likelihoods
    shift = torch.zeros_like(ca,device=device,requires_grad=True)
    if steps>0:
        minimizer = torch.optim.Adam([shift], lr=0.5)
        for i in range(steps):
            minimizer.zero_grad()
            nll1 = nll(ca+shift)
            nll1.backward()
            minimizer.step()
    
    return ca+shift
//  [python]
def init_xyz_edm(mu,sigma,batch=1,steps=50):
    '''generate initial Ca trace from mu,sigma of 
    predicted log-normal distributions'''
    device = mu.device
    L = mu.shape[-1]
    m = mu.detach()
    s = sigma.detach()
    
    # approximate initial distances by the mode
    # of lognormal distributions
    D = torch.exp(m-s*s)
    sampler = LogNormal(mu[0],sigma[0])
    Ds = sampler.sample((batch,))
    Ds[0] = D[0]
    
    # convert distance matrix to a Gramm matrix
    J = torch.eye(L,device=device)-torch.ones((L,L),device=device)/L
    G = -0.5*J@Ds**2@J
    
    # use three leading eigenvectors to approximate
    # initial coordinates
    e,v = torch.linalg.eigh(G)
    ca = v[:,:,-3:]*e[:,-3:].unsqueeze(1)**0.5
    
    def nll(xyz):
        dist = torch.cdist(xyz,xyz)+1e-6*torch.eye(L,device=device)
        nll = (0.5*((torch.log(dist)-m)/s + torch.log(dist*s*(2*np.pi)**0.5))**2).mean()
        return nll
    
    nll0 = nll(ca)
    nll1 = nll0
    
    # fine-tuning by minimization of log-likelihoods
    shift = torch.zeros_like(ca,device=device,requires_grad=True)
    if steps>0:
        minimizer = torch.optim.Adam([shift], lr=0.5)
        for i in range(steps):
            minimizer.zero_grad()
            nll1 = nll(ca+shift)
            nll1.backward()
            minimizer.step()
    
    return ca+shift
//  [python]
def train(self, x):
        '''training pca.
        Args:
            x: [N, dim] FloatTensor containing data which undergoes PCA/Whitening.
        '''

        print('calculating PCA parameters...')

        ### Referring to https://github.com/Relja/relja_matlab/blob/master/relja_PCA.m

        x = x.t()
        nPoints = x.size(1)
        nDims = x.size(0)

        # x = x.double()
        mu = x.mean(1).unsqueeze(1)
        x = x - mu

        if (nDims<=nPoints):
            doDual = False
            x2 = torch.matmul(x, x.t()) / (nPoints - 1)
        else:
            doDual = True
            x2 = torch.matmul(x.t(), x) / (nPoints - 1)

        L, U = torch.symeig(x2, eigenvectors=True)
        if (self.pca_n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.pca_n_components]
            L = torch.index_select(L, 0, k_indices)
            U = torch.index_select(U, 1, k_indices)

        lams = L
        lams[lams<1e-9] = 1e-9

        if (doDual):
            U = torch.matmul(x, torch.matmul(U, torch.diag(1./torch.sqrt(lams))/np.sqrt(nPoints-1)))

        Utmu= torch.matmul(U.t(), mu)

        U, lams, mu, Utmu = U.numpy(), lams.numpy(), mu.numpy(), Utmu.numpy()

        # save as h5 file.
        print('================= PCA RESULT ==================')
        print('U: {}'.format(U.shape))
        print('lams: {}'.format(lams.shape))
        print('mu: {}'.format(mu.shape))
        print('Utmu: {}'.format(Utmu.shape))
        print('===============================================')

        # save features, labels to h5 file.
        filename = os.path.join(self.pca_parameters_path)
        h5file = h5py.File(filename, 'w')
        h5file.create_dataset('U', data=U)
        h5file.create_dataset('lams', data=lams)
        h5file.create_dataset('mu', data=mu)
        h5file.create_dataset('Utmu', data=Utmu)
        h5file.close()
//  [python]
def train(self, x):
        '''training pca.
        Args:
            x: [N, dim] FloatTensor containing data which undergoes PCA/Whitening.
        '''

        print('calculating PCA parameters...')

        ### Referring to https://github.com/Relja/relja_matlab/blob/master/relja_PCA.m

        x = x.t()
        nPoints = x.size(1)
        nDims = x.size(0)

        # x = x.double()
        mu = x.mean(1).unsqueeze(1)
        x = x - mu

        if (nDims<=nPoints):
            doDual = False
            x2 = torch.matmul(x, x.t()) / (nPoints - 1)
        else:
            doDual = True
            x2 = torch.matmul(x.t(), x) / (nPoints - 1)

        L, U = torch.linalg.eigh(x2)
        if (self.pca_n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.pca_n_components]
            L = torch.index_select(L, 0, k_indices)
            U = torch.index_select(U, 1, k_indices)

        lams = L
        lams[lams<1e-9] = 1e-9

        if (doDual):
            U = torch.matmul(x, torch.matmul(U, torch.diag(1./torch.sqrt(lams))/np.sqrt(nPoints-1)))

        Utmu= torch.matmul(U.t(), mu)

        U, lams, mu, Utmu = U.numpy(), lams.numpy(), mu.numpy(), Utmu.numpy()

        # save as h5 file.
        print('================= PCA RESULT ==================')
        print('U: {}'.format(U.shape))
        print('lams: {}'.format(lams.shape))
        print('mu: {}'.format(mu.shape))
        print('Utmu: {}'.format(Utmu.shape))
        print('===============================================')

        # save features, labels to h5 file.
        filename = os.path.join(self.pca_parameters_path)
        h5file = h5py.File(filename, 'w')
        h5file.create_dataset('U', data=U)
        h5file.create_dataset('lams', data=lams)
        h5file.create_dataset('mu', data=mu)
        h5file.create_dataset('Utmu', data=Utmu)
        h5file.close()
//  [python]
def get_frame(self,pnts,override=None):
        if override is None:
            is_local_frame = self.is_local_frame
        else:
            is_local_frame = override
        if is_local_frame:
            batch_size = pnts.size(0)
            num_points = pnts.size(1)
            
            def knn(x, k):
               
                x_squrae = (x**2).sum(-1,True)
                pairwise_distance = -(x_squrae -2*torch.bmm(x,x.transpose(1,2))  + x_squrae.transpose(1,2))
            
                idx = pairwise_distance.topk(k=k, dim=-1)[1]   # (batch_size, num_points, k)
                return idx
            idx = knn(pnts, k=self.k_size)
            device = torch.device('cuda')

            idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points

            idx = idx + idx_base

            idx = idx.view(-1)
             
            pnts = pnts.view(batch_size*num_points, -1)[idx, :].view(batch_size, num_points, self.k_size, 3) 
                
            center = pnts.mean(2,False)
            pnts_centered = pnts - center.unsqueeze(2)
            
            R =  torch.einsum('bpki,bpkj->bpij',pnts_centered,pnts_centered)
            lambdas,V_ = torch.symeig(R.detach().cpu(),True)            
            F =  V_.to(R)
            if self.is_detach_frame:
                return F.detach(), center.detach(),pnts_centered.detach()
            else:
                return F,center,pnts_centered

        else:
            center = pnts.mean(1,True)
            pnts_centered = pnts - center
            R = torch.bmm(pnts_centered.transpose(1,2),pnts_centered)
            lambdas,V_ = torch.symeig(R.detach().cpu(),True)            
            F =  V_.to(R).unsqueeze(1).repeat(1,pnts.shape[1],1,1)
            if self.is_detach_frame:
                return F.detach(), center.detach()
            else:
                return F,center
//  [python]
def get_frame(self,pnts,override=None):
        if override is None:
            is_local_frame = self.is_local_frame
        else:
            is_local_frame = override
        if is_local_frame:
            batch_size = pnts.size(0)
            num_points = pnts.size(1)
            
            def knn(x, k):
               
                x_squrae = (x**2).sum(-1,True)
                pairwise_distance = -(x_squrae -2*torch.bmm(x,x.transpose(1,2))  + x_squrae.transpose(1,2))
            
                idx = pairwise_distance.topk(k=k, dim=-1)[1]   # (batch_size, num_points, k)
                return idx
            idx = knn(pnts, k=self.k_size)
            device = torch.device('cuda')

            idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points

            idx = idx + idx_base

            idx = idx.view(-1)
             
            pnts = pnts.view(batch_size*num_points, -1)[idx, :].view(batch_size, num_points, self.k_size, 3) 
                
            center = pnts.mean(2,False)
            pnts_centered = pnts - center.unsqueeze(2)
            
            R =  torch.einsum('bpki,bpkj->bpij',pnts_centered,pnts_centered)
            lambdas,V_ = torch.linalg.eigh(R.detach().cpu())            
            F =  V_.to(R)
            if self.is_detach_frame:
                return F.detach(), center.detach(),pnts_centered.detach()
            else:
                return F,center,pnts_centered

        else:
            center = pnts.mean(1,True)
            pnts_centered = pnts - center
            R = torch.bmm(pnts_centered.transpose(1,2),pnts_centered)
            lambdas,V_ = torch.linalg.eigh(R.detach().cpu())            
            F =  V_.to(R).unsqueeze(1).repeat(1,pnts.shape[1],1,1)
            if self.is_detach_frame:
                return F.detach(), center.detach()
            else:
                return F,center
//  [python]
def stn(self, x, u):
        # A_vec = self.fc_stn(u)
        A_vec = self.fc_stn(torch.cat([u, x.reshape(-1, 28 * 28)], 1))
        A = convert_Avec_to_A(A_vec)
        _, evs = torch.symeig(A, eigenvectors=True)
        tcos, tsin = evs[:, 0:1, 0:1], evs[:, 1:2, 0:1]

        self.theta_angle = torch.atan2(tsin[:, 0, 0], tcos[:, 0, 0])

        # clock-wise rotate theta
        theta_0 = torch.cat([tcos, tsin, tcos * 0], 2)
        theta_1 = torch.cat([-tsin, tcos, tcos * 0], 2)
        theta = torch.cat([theta_0, theta_1], 1)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)
        return x
//  [python]
def stn(self, x, u):
        A_vec = self.fc_stn(torch.cat([u, x.reshape(-1, 28 * 28)], 1))
        A = convert_Avec_to_A(A_vec)
        _, evs = torch.linalg.eigh(A)
        tcos, tsin = evs[:, 0:1, 0:1], evs[:, 1:2, 0:1]

        self.theta_angle = torch.atan2(tsin[:, 0, 0], tcos[:, 0, 0])

        theta_0 = torch.cat([tcos, tsin, tcos * 0], 2)
        theta_1 = torch.cat([-tsin, tcos, tcos * 0], 2)
        theta = torch.cat([theta_0, theta_1], 1)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)
        return x
//  [python]
def deterministic_quadrature(T, dim):
    eigvals, _ = torch.symeig(T, eigenvectors=True)
    clamped_eigvals = eigvals.clamp(min=0).detach() + (eigvals - eigvals.detach())
    return torch.mean(torch.log(clamped_eigvals + 1e-8), dim=-1) * dim
//  [python]
def deterministic_quadrature(T, dim):
    eigvals, _ = torch.linalg.eigh(T)
    clamped_eigvals = eigvals.clamp(min=0).detach() + (eigvals - eigvals.detach())
    return torch.mean(torch.log(clamped_eigvals + 1e-8), dim=-1) * dim
//  [python]
def stochastic_trfA(mvp_fun, v, m):
    _, dim = v.shape
    T = lanczos_tridiagonalization(mvp_fun, m, v)
    T_eigvals, T_eigvecs = torch.symeig(T, eigenvectors=True)
    clamped_eigvals = T_eigvals.clamp(min=0).detach() + (T_eigvals - T_eigvals.detach())
    tau = T_eigvecs[..., 0, :]
    estimates = torch.sum(tau**2 * torch.log(clamped_eigvals + 1e-8), dim=-1) * dim
    return estimates
//  [python]
def stochastic_trfA(mvp_fun, v, m):
    _, dim = v.shape
    T = lanczos_tridiagonalization(mvp_fun, m, v)
    T_eigvals, T_eigvecs = torch.linalg.eigh(T)
    clamped_eigvals = T_eigvals.clamp(min=0).detach() + (T_eigvals - T_eigvals.detach())
    tau = T_eigvecs[..., 0, :]
    estimates = torch.sum(tau**2 * torch.log(clamped_eigvals + 1e-8), dim=-1) * dim
    return estimates
//  [python]
def stochastic_quadrature(T, dim, func=torch.log):
    eigvals, eigvecs = torch.symeig(T, eigenvectors=True)

    with torch.no_grad():
        if eigvals.numel() > torch.unique(eigvals, dim=1).numel():
            print("Non-unique eigenvalues.")

    clamped_eigvals = eigvals.clamp(min=0).detach() + (eigvals - eigvals.detach())
    tau = eigvecs[..., 0, :]
    return torch.sum(tau * tau * func(clamped_eigvals + 1e-8), dim=-1) * dim
//  [python]
def stochastic_quadrature(T, dim, func=torch.log):
    eigvals, eigvecs = torch.linalg.eigh(T)

    with torch.no_grad():
        if eigvals.numel() > torch.unique(eigvals, dim=1).numel():
            print("Non-unique eigenvalues.")

    clamped_eigvals = eigvals.clamp(min=0).detach() + (eigvals - eigvals.detach())
    tau = eigvecs[..., 0, :]
    return torch.sum(tau * tau * func(clamped_eigvals + 1e-8), dim=-1) * dim
//  [python]
def FISTA(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver FISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
        ahat_y = ahat_y.add(eta * basis.t().mm(Res))
        ahat = ahat_y.sub(eta * lambd).clamp(min = 0.)
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def FISTA(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver FISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()))[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
        ahat_y = ahat_y.add(eta * basis.t().mm(Res))
        ahat = ahat_y.sub(eta * lambd).clamp(min = 0.)
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def compute_ntk_trace(net, inputs, targets, mode, loss_fn, split_data=1):
    # Compute gradients (but don't apply them)
    net.zero_grad()
    N = inputs.shape[0]
    grads = []

    autograd_hacks.add_hooks(net)
    outputs = net.forward(inputs)
    sum(outputs[torch.arange(N), targets]).backward()
    autograd_hacks.compute_grad1(net, loss_type='sum')

    grads = [
        param.grad1.flatten(start_dim=1) for param in net.parameters()
        if hasattr(param, 'grad1')
    ]
    grads = torch.cat(grads, axis=1)

    ntk = torch.matmul(grads, grads.t())
    eigenvalues, _ = torch.symeig(ntk)  # ascending
    return np.sqrt(eigenvalues.cpu().numpy().sum() / N)
//  [python]
def compute_ntk_trace(net, inputs, targets, mode, loss_fn, split_data=1):
    # Compute gradients (but don't apply them)
    net.zero_grad()
    N = inputs.shape[0]
    grads = []

    autograd_hacks.add_hooks(net)
    outputs = net.forward(inputs)
    sum(outputs[torch.arange(N), targets]).backward()
    autograd_hacks.compute_grad1(net, loss_type='sum')

    grads = [
        param.grad1.flatten(start_dim=1) for param in net.parameters()
        if hasattr(param, 'grad1')
    ]
    grads = torch.cat(grads, axis=1)

    ntk = torch.matmul(grads, grads.t())
    eigenvalues = torch.linalg.eigh(ntk).eigenvalues  # ascending
    return np.sqrt(eigenvalues.cpu().numpy().sum() / N)
//  [python]
def compute_eig(self,
                    d_mat,
                    train):

        dim = d_mat.shape[-1]
        # do analytically if possible to avoid sign ambiguity
        # in the eigenvectors, which leads to worse training
        # results for the nacv
        if dim == 2:
            ad_energies, u = self.diag(d_mat)
        # otherwise do numerically
        else:
            # catch any nans - if you leave them before
            # calculating the eigenvectors then they will
            # raise an error
            nan_idx = torch.bitwise_not(torch.isfinite(d_mat)).any(-1).any(-1)

            if not train:
                d_mat[nan_idx, :, :] = 0

            # diagonalize
            ad_energies, u = torch.symeig(d_mat, True)

        return ad_energies, u, nan_idx
//  [python]
def compute_eig(self,
                    d_mat,
                    train):

        dim = d_mat.shape[-1]
        # do analytically if possible to avoid sign ambiguity
        # in the eigenvectors, which leads to worse training
        # results for the nacv
        if dim == 2:
            ad_energies, u = self.diag(d_mat)
        # otherwise do numerically
        else:
            # catch any nans - if you leave them before
            # calculating the eigenvectors then they will
            # raise an error
            nan_idx = torch.bitwise_not(torch.isfinite(d_mat)).any(-1).any(-1)

            if not train:
                d_mat[nan_idx, :, :] = 0

            # diagonalize
            ad_energies, u = torch.linalg.eigh(d_mat)

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

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

            self._inversed = True
            self.iter += 1
//  [python]
def nacv_from_hess(self,
                       batch,
                       results,
                       states):

        gaps, gap_hess = self.process_hess(batch=batch,
                                           results=results,
                                           states=states)
        nacvs = []

        for gap, hess in zip(gaps, gap_hess):

            # Return 0 if the gap is > threshold
            # This is very inefficient right now because the
            # Hessian is computed even if it doesn't get used

            dim = int((hess.reshape(-1).shape[0]) ** 0.5)

            if gap > self.gap_threshold:
                nacv = (torch.zeros(dim // 3, 3)
                        .to(gap.device))
                nacvs.append(nacv)
                continue

            # if you simplify eqs. 5 and 6 in the SchNarc
            # paper you get this

            nacv_kron = (1 / 4 * gap * hess).reshape(dim, dim)

            eigs, vecs = torch.symeig(nacv_kron, True)
            argmax = eigs.argmax()
            lam = eigs[argmax]
            v = vecs[:, argmax]
            nacv = v * (lam ** 0.5)

            nacvs.append(nacv.reshape(-1, 3))

        nacvs = torch.cat(nacvs)

        return nacvs
//  [python]
def nacv_from_hess(self,
                       batch,
                       results,
                       states):

        gaps, gap_hess = self.process_hess(batch=batch,
                                           results=results,
                                           states=states)
        nacvs = []

        for gap, hess in zip(gaps, gap_hess):

            # Return 0 if the gap is > threshold
            # This is very inefficient right now because the
            # Hessian is computed even if it doesn't get used

            dim = int((hess.reshape(-1).shape[0]) ** 0.5)

            if gap > self.gap_threshold:
                nacv = (torch.zeros(dim // 3, 3)
                        .to(gap.device))
                nacvs.append(nacv)
                continue

            # if you simplify eqs. 5 and 6 in the SchNarc
            # paper you get this

            nacv_kron = (1 / 4 * gap * hess).reshape(dim, dim)

            eigs, vecs = torch.linalg.eigh(nacv_kron)
            argmax = eigs.argmax()
            lam = eigs[argmax]
            v = vecs[:, argmax]
            nacv = v * (lam ** 0.5)

            nacvs.append(nacv.reshape(-1, 3))

        nacvs = torch.cat(nacvs)

        return nacvs
//  [python]
def test_symeig(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            oute = torch.empty(dims[1:] + dims[:1], dtype=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:
                x_recon = torch.matmul(torch.matmul(outv, torch.diag_embed(oute)), outv.transpose(-2, -1))
                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_symmetric_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:
                x_recon = torch.matmul(torch.matmul(resv, torch.diag_embed(rese)), resv.transpose(-2, -1))
                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 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_symmetric_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            oute = torch.empty(dims[1:] + dims[:1], dtype=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:
                x_recon = torch.matmul(torch.matmul(outv, torch.diag_embed(oute)), outv.transpose(-2, -1))
                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_symmetric_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:
                x_recon = torch.matmul(torch.matmul(resv, torch.diag_embed(rese)), resv.transpose(-2, -1))
                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 product(batch_dims_set, (True, False), (True, False)):
            run_test((5,) + batch_dims, eigenvectors, upper)
//  [python]
def lanczos_approx(self, L, order, e1, dx, u):
        v, H_M = self.planczos(L, order, dx)
        H_M_eval, H_M_evec = torch.symeig(H_M, eigenvectors=True)
        H_M_eval = torch.clamp(H_M_eval, 0, H_M_eval.max().item())
        fv = H_M_evec @ torch.diag_embed(f(H_M_eval, u)) @ H_M_evec.permute(0, 1, 3, 2)
        approx = torch.norm(dx, dim=2).unsqueeze(-1).unsqueeze(-1) * v @ fv @ e1
        return approx
//  [python]
def lanczos_approx(self, L, order, e1, dx, u):
        v, H_M = self.planczos(L, order, dx)
        H_M_eval, H_M_evec = torch.linalg.eigh(H_M)
        H_M_eval = torch.clamp(H_M_eval, 0, H_M_eval.max().item())
        fv = H_M_evec @ torch.diag_embed(f(H_M_eval, u)) @ H_M_evec.permute(0, 1, 3, 2)
        approx = torch.norm(dx, dim=2).unsqueeze(-1).unsqueeze(-1) * v @ fv @ e1
        return approx
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor(
            [[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0],
             [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
             [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.symeig(K, True, False)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor(
            [[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0],
             [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
             [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.linalg.eigh(K)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
	# I referenced the default svd tolerance value in matlab.

	b, c, h, w = feature.size()

	mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
	zeromean = feature.view(b, c, -1) - mean
	cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
	
	eps_matrix_ = (torch.ones_like(cov)*1e-8).cuda()
	#evals, evects = torch.symeig(cov, eigenvectors=True)
	evals, evects = torch.symeig(cov+eps_matrix_, eigenvectors=True)
	
	
	p = 0.5
	if inverse:
		p *= -1

	covsqrt = []
	for i in range(b):
		k = 0
		for j in range(c):
			if evals[i][j] > tolerance:
				k = j
				break
		covsqrt.append(torch.mm(evects[i][:, k:],
							torch.mm(evals[i][k:].pow(p).diag_embed(),
									 evects[i][:, k:].t())).unsqueeze(0))
	covsqrt = torch.cat(covsqrt, dim=0)

	return covsqrt, mean
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
	# I referenced the default svd tolerance value in matlab.

	b, c, h, w = feature.size()

	mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
	zeromean = feature.view(b, c, -1) - mean
	cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
	
	eps_matrix_ = (torch.ones_like(cov)*1e-8).cuda()
	evals, evects = torch.linalg.eigh(cov+eps_matrix_)
	
	p = 0.5
	if inverse:
		p *= -1

	covsqrt = []
	for i in range(b):
		k = 0
		for j in range(c):
			if evals[i][j] > tolerance:
				k = j
				break
		covsqrt.append(torch.mm(evects[i][:, k:],
							torch.mm(evals[i][k:].pow(p).diag_embed(),
									 evects[i][:, k:].t())).unsqueeze(0))
	covsqrt = torch.cat(covsqrt, dim=0)

	return covsqrt, mean
//  [python]
def get_ntk_n(xloader, networks, recalbn=0, train_mode=False, num_batch=-1):
    device = 'cpu'
    ntks = []
    
    for network in networks:
        #network.to(device)
        if train_mode:
            network.train()
        else:
            network.eval()
    ######
    grads = [[] for _ in range(len(networks))]
    for i, (inputs, targets) in enumerate(xloader):
        if num_batch > 0 and i >= num_batch: break
        inputs = inputs.to(device)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone().to(device=device, non_blocking=True)
            logit = network(inputs_)
            if isinstance(logit, tuple):
                logit = logit[1]  # 201 networks: return features and logits
            for _idx in range(len(inputs_)):
                logit[_idx:_idx+1].backward(torch.ones_like(logit[_idx:_idx+1]), retain_graph=True)
                grad = []
                for name, W in network.named_parameters():
                    if 'weight' in name and W.grad is not None:
                        grad.append(W.grad.view(-1).detach())
                grads[net_idx].append(torch.cat(grad, -1))
                network.zero_grad()
                torch.cuda.empty_cache()
    ######
    grads = [torch.stack(_grads, 0) for _grads in grads]
    ntks = [torch.einsum('nc,mc->nm', [_grads, _grads]) for _grads in grads]
    conds = []
    for ntk in ntks:
        eigenvalues, _ = torch.symeig(ntk)  # ascending
        conds.append(np.nan_to_num((eigenvalues[-1] / eigenvalues[0]).item(), copy=True, nan=100000.0))
    return conds
//  [python]
def get_ntk_n(xloader, networks, recalbn=0, train_mode=False, num_batch=-1):
    device = 'cpu'
    ntks = []
    
    for network in networks:
        #network.to(device)
        if train_mode:
            network.train()
        else:
            network.eval()
    ######
    grads = [[] for _ in range(len(networks))]
    for i, (inputs, targets) in enumerate(xloader):
        if num_batch > 0 and i >= num_batch: break
        inputs = inputs.to(device)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone().to(device=device, non_blocking=True)
            logit = network(inputs_)
            if isinstance(logit, tuple):
                logit = logit[1]  # 201 networks: return features and logits
            for _idx in range(len(inputs_)):
                logit[_idx:_idx+1].backward(torch.ones_like(logit[_idx:_idx+1]), retain_graph=True)
                grad = []
                for name, W in network.named_parameters():
                    if 'weight' in name and W.grad is not None:
                        grad.append(W.grad.view(-1).detach())
                grads[net_idx].append(torch.cat(grad, -1))
                network.zero_grad()
                torch.cuda.empty_cache()
    ######
    grads = [torch.stack(_grads, 0) for _grads in grads]
    ntks = [torch.einsum('nc,mc->nm', [_grads, _grads]) for _grads in grads]
    conds = []
    for ntk in ntks:
        eigenvalues = torch.linalg.eigh(ntk).eigenvalues  # ascending
        conds.append(np.nan_to_num((eigenvalues[-1] / eigenvalues[0]).item(), copy=True, nan=100000.0))
    return conds
//  [python]
def loss(self, H1, H2):
        """
        It is the loss function of CCA as introduced in the original paper. There can be other formulations.
        """

        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()

        o1 =  H1.size(0)
        o2 = H2.size(0)

        m = H1.size(1)



        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=self.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=self.device)

        # Workaround !!! USE BATCH > 16
        [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
        [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)
        
        # Newest way but not debugged
        # [D1, V1] = torch.linalg.eig(SigmaHat11)
        # [D2, V2] = torch.linalg.eig(SigmaHat22)
        
        # Original
        # [D1, V1] = torch.eig(SigmaHat11, eigenvectors=True)
        # [D2, V2] = torch.eig(SigmaHat22, eigenvectors=True)

        D1 = D1.unsqueeze(1)
        posInd1 = torch.gt(D1[:, 0], eps).nonzero()[:, 0]
        D1 = D1[posInd1, 0]
        V1 = V1[:, posInd1]
        D1 = torch.squeeze(D1)  # Remove extra dimensions from D1
        # Reshape D1 to 1-dimensional tensor
        D2 = D2.unsqueeze(1)
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        posInd2 = torch.gt(D2[:, 0], eps).nonzero()[:, 0]
        D2 = D2[posInd2, 0]
        V2 = V2[:, posInd2]
        D2 = torch.squeeze(D2)
        
        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)


        if self.use_all_singular_values:
    
            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))
  
        else:
   
            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(self.device)) 
            U, V = torch.symeig(trace_TT, eigenvectors=True)
            U = torch.where(U>eps, U, (torch.ones(U.shape).float()*eps).to(self.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def loss(self, H1, H2):
        """
        It is the loss function of CCA as introduced in the original paper. There can be other formulations.
        """

        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()

        o1 =  H1.size(0)
        o2 = H2.size(0)

        m = H1.size(1)



        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=self.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=self.device)

        # Workaround !!! USE BATCH > 16
        [D1, V1] = torch.linalg.eigh(SigmaHat11)
        [D2, V2] = torch.linalg.eigh(SigmaHat22)
        
        D1 = D1.unsqueeze(1)
        posInd1 = torch.gt(D1[:, 0], eps).nonzero()[:, 0]
        D1 = D1[posInd1, 0]
        V1 = V1[:, posInd1]
        D1 = torch.squeeze(D1)  # Remove extra dimensions from D1
        # Reshape D1 to 1-dimensional tensor
        D2 = D2.unsqueeze(1)
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        posInd2 = torch.gt(D2[:, 0], eps).nonzero()[:, 0]
        D2 = D2[posInd2, 0]
        V2 = V2[:, posInd2]
        D2 = torch.squeeze(D2)
        
        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)


        if self.use_all_singular_values:
    
            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))
  
        else:
   
            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(self.device)) 
            U, V = torch.linalg.eigh(trace_TT)
            U = torch.where(U>eps, U, (torch.ones(U.shape).float()*eps).to(self.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr