jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def LDAloss(H, label, device, data_type):
    N, C, d, _ = H.shape
    N_new = N * d * d
    H = H.permute(0, 2, 3, 1)
    H = torch.reshape(H, (N_new, C))
    H_bar = H - torch.mean(H, 0, True)
    label = label.view(N, 1)
    labels = torch.reshape(label * torch.Tensor().new_ones((N, d * d), device=device, dtype=torch.long), (N_new,))
    S_w = torch.Tensor().new_zeros((C, C), device=device, dtype=data_type)
    S_t = H_bar.t().matmul(H_bar) / (N_new - 1)
    for i in range(categ):
        H_i = H[torch.nonzero(labels == i).view(-1)]
        H_i_bar = H_i - torch.mean(H_i, 0, True)
        N_i = H_i.shape[0]
        if N_i == 0:
            continue
        S_w += H_i_bar.t().matmul(H_i_bar) / (N_i - 1) / categ
    temp = (S_w + lamb * torch.diag(torch.Tensor().new_ones((C), device=device, dtype=data_type))).pinverse().matmul(
        S_t - S_w)
    w, v = torch.symeig(temp, eigenvectors=True)
    w = w.detach()
    v = v.detach()
    loss = (v[:, 12 - eig_num: 12].t().matmul(temp).matmul(v[:, 12 - eig_num: 12])).sum()
    return -loss / eig_num, w
//  [python]
def LDAloss(H, label, device, data_type):
    N, C, d, _ = H.shape
    N_new = N * d * d
    H = H.permute(0, 2, 3, 1)
    H = torch.reshape(H, (N_new, C))
    H_bar = H - torch.mean(H, 0, True)
    label = label.view(N, 1)
    labels = torch.reshape(label * torch.Tensor().new_ones((N, d * d), device=device, dtype=torch.long), (N_new,))
    S_w = torch.Tensor().new_zeros((C, C), device=device, dtype=data_type)
    S_t = H_bar.t().matmul(H_bar) / (N_new - 1)
    for i in range(categ):
        H_i = H[torch.nonzero(labels == i).view(-1)]
        H_i_bar = H_i - torch.mean(H_i, 0, True)
        N_i = H_i.shape[0]
        if N_i == 0:
            continue
        S_w += H_i_bar.t().matmul(H_i_bar) / (N_i - 1) / categ
    temp = (S_w + lamb * torch.diag(torch.Tensor().new_ones((C), device=device, dtype=data_type))).pinverse().matmul(
        S_t - S_w)
    w, v = torch.linalg.eigh(temp)
    w = w.detach()
    v = v.detach()
    loss = (v[:, 12 - eig_num: 12].t().matmul(temp).matmul(v[:, 12 - eig_num: 12])).sum()
    return -loss / eig_num, w
//  [python]
def getSHOTLocalRF(centralPoints, neighs, radius, nnidx):
    displacement = (neighs - centralPoints.unsqueeze(2))                     # (B, N', k, 3)
    dists = torch.sqrt(torch.sum(torch.pow(displacement, 2), dim=-1))        # (B, N', k)

    w = (radius - dists).clamp(min=0.0)                                      # (B, N', k)

    # idx mask
    idx_mask = (nnidx - nnidx[:,:,0].unsqueeze(-1)).bool()
    idx_mask[:,:,0] = True 
    idx_bool = torch.sum(nnidx, dim=-1).bool()                               # for all nnidx are 0
    idx_mask[~idx_bool] = False

    w = w * idx_mask
    sumw = torch.sum(w, dim=-1, keepdim=True)                                # (B, N', 1)

    weightDis = torch.einsum('bnkc,bnk->bnkc', displacement, torch.sqrt(w))  # (B, N', k, 3)

    M = torch.einsum('bnkc,bnkd->bncd', weightDis, weightDis)                # (B, N', 3, 3)
    M = M / (sumw.unsqueeze(2) + 1e-10)

    # Eigenvalue decomposition
    eVal, eVec = torch.symeig(M.cpu(), eigenvectors=True)                    # (B, N', 3) (B, N', 3, 3) ascending order
    # assert((eVal[:,:,0]<=eVal[:,:,1]).all() and (eVal[:,:,1]<=eVal[:,:,2]).all()), "[ERROR] eigenvalues should be ascending order."

    ref_Z = eVec[:,:,:,0].cuda()  # min  (B, N', 3)
    ref_X = eVec[:,:,:,2].cuda()  # max  (B, N', 3)

    # Disambiguate x and z
    xInFeatRef = torch.einsum('bnkc,bnc->bnk', displacement, ref_X)          # (B, N', k)
    zInFeatRef = torch.einsum('bnkc,bnc->bnk', displacement, ref_Z)          # (B, N', k)

    disX_pos = ((xInFeatRef * idx_mask) > 0) 
    disX_pos = torch.sum(disX_pos, dim=-1, keepdim=True)                     # (B, N', 1)
    disZ_pos = ((zInFeatRef * idx_mask) > 0) 
    disZ_pos = torch.sum(disZ_pos, dim=-1, keepdim=True)                     # (B, N', 1)

    nk = torch.sum(idx_mask, dim=-1, keepdim=True)
    ref_X = torch.where(disX_pos < (nk-disX_pos), -ref_X, ref_X)
    ref_Z = torch.where(disZ_pos < (nk-disZ_pos), -ref_Z, ref_Z)

    ref_Y = torch.cross(ref_Z, ref_X, dim=-1)
    
    return ref_X, ref_Y, ref_Z
//  [python]
def getSHOTLocalRF(centralPoints, neighs, radius, nnidx):
    displacement = (neighs - centralPoints.unsqueeze(2))                     # (B, N', k, 3)
    dists = torch.sqrt(torch.sum(torch.pow(displacement, 2), dim=-1))        # (B, N', k)

    w = (radius - dists).clamp(min=0.0)                                      # (B, N', k)

    # idx mask
    idx_mask = (nnidx - nnidx[:,:,0].unsqueeze(-1)).bool()
    idx_mask[:,:,0] = True 
    idx_bool = torch.sum(nnidx, dim=-1).bool()                               # for all nnidx are 0
    idx_mask[~idx_bool] = False

    w = w * idx_mask
    sumw = torch.sum(w, dim=-1, keepdim=True)                                # (B, N', 1)

    weightDis = torch.einsum('bnkc,bnk->bnkc', displacement, torch.sqrt(w))  # (B, N', k, 3)

    M = torch.einsum('bnkc,bnkd->bncd', weightDis, weightDis)                # (B, N', 3, 3)
    M = M / (sumw.unsqueeze(2) + 1e-10)

    # Eigenvalue decomposition
    eVal, eVec = torch.linalg.eigh(M)                    # (B, N', 3) (B, N', 3, 3) ascending order
    # assert((eVal[:,:,0]<=eVal[:,:,1]).all() and (eVal[:,:,1]<=eVal[:,:,2]).all()), "[ERROR] eigenvalues should be ascending order."

    ref_Z = eVec[:,:,:,0]  # min  (B, N', 3)
    ref_X = eVec[:,:,:,2]  # max  (B, N', 3)

    # Disambiguate x and z
    xInFeatRef = torch.einsum('bnkc,bnc->bnk', displacement, ref_X)          # (B, N', k)
    zInFeatRef = torch.einsum('bnkc,bnc->bnk', displacement, ref_Z)          # (B, N', k)

    disX_pos = ((xInFeatRef * idx_mask) > 0) 
    disX_pos = torch.sum(disX_pos, dim=-1, keepdim=True)                     # (B, N', 1)
    disZ_pos = ((zInFeatRef * idx_mask) > 0) 
    disZ_pos = torch.sum(disZ_pos, dim=-1, keepdim=True)                     # (B, N', 1)

    nk = torch.sum(idx_mask, dim=-1, keepdim=True)
    ref_X = torch.where(disX_pos < (nk-disX_pos), -ref_X, ref_X)
    ref_Z = torch.where(disZ_pos < (nk-disZ_pos), -ref_Z, ref_Z)

    ref_Y = torch.cross(ref_Z, ref_X, dim=-1)
    
    return ref_X, ref_Y, ref_Z
//  [python]
def eigenvalues(self):
        eig_A, _ = torch.symeig(self.A)
        eig_B, _ = torch.symeig(self.B)
        eig = torch.ger(eig_A, eig_B).flatten()
        return torch.sort(eig, descending=True)[0]
//  [python]
def eigenvalues(self):
        eig_A, _ = torch.linalg.eigh(self.A)
        eig_B, _ = torch.linalg.eigh(self.B)
        eig = torch.ger(eig_A, eig_B).flatten()
        return torch.sort(eig, descending=True)[0]
//  [python]
def loss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

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

    o1 = o2 = H1.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=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

    # Calculating the root inverse of covariance matrices by using eigen decomposition
    [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
    [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

    # Added to increase stability
    posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
    D1 = D1[posInd1]
    V1 = V1[:, posInd1]
    posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
    D2 = D2[posInd2]
    V2 = V2[:, posInd2]

    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 use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.symeig(trace_TT, eigenvectors=True)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def loss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

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

    o1 = o2 = H1.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=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

    # Calculating the root inverse of covariance matrices by using eigen decomposition
    [D1, V1] = torch.linalg.eigh(SigmaHat11)
    [D2, V2] = torch.linalg.eigh(SigmaHat22)

    # Added to increase stability
    posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
    D1 = D1[posInd1]
    V1 = V1[:, posInd1]
    posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
    D2 = D2[posInd2]
    V2 = V2[:, posInd2]

    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 use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.linalg.eigh(trace_TT)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def top_eigenvalue(self):
        eig_A, _ = torch.symeig(self.A)
        eig_B, _ = torch.symeig(self.B)
        return (eig_A.max() * eig_B.max()).item()
//  [python]
def top_eigenvalue(self):
        eig_A, _ = torch.linalg.eigh(self.A)
        eig_B, _ = torch.linalg.eigh(self.B)
        return (eig_A.max() * eig_B.max()).item()
//  [python]
def solve_tica_eigenproblem(self,C_0,C_lag,n_eig=0,save_params=False):
        '''
        Compute generalized eigenvalue problem : C_lag * wi = lambda_i * C_0 * w_i

        Parameters:
            C_0 (tensor): correlation matrix at lag-time 0
            C_lag (tensor): correlation matrix at lag-time lag
            scipy (bool): wheter to use scipy.eigh to compute eigenvecs (debug only)

        Returns:
            eigvals (tensor): eigenvalues in ascending order
            #eigvecs (tensor): matrix whose column eigvecs[:,i] is eigenvector associated to eigvals[i]
        '''
        #cholesky decomposition
        if self.reg_cholesky is not None:
            L = torch.cholesky(C_0+self.reg_cholesky*torch.eye(C_0.shape[0]),upper=False)
        else:
            L = torch.cholesky(C_0,upper=False)
            
        L_t = torch.t(L)
        L_ti = torch.inverse(L_t)
        L_i = torch.inverse(L)
        C_new = torch.matmul(torch.matmul(L_i,C_lag),L_ti)

        #find eigenvalues and vectors of C_new
        eigvals, eigvecs = torch.symeig(C_new,eigenvectors=True)
        #sort
        eigvals, indices = torch.sort(eigvals, 0, descending=True)
        eigvecs = eigvecs[:,indices]
        
        #return to original eigenvectors
        eigvecs = torch.matmul(L_ti,eigvecs)

        #normalize them
        for i in range(eigvecs.shape[1]): # maybe change in sum along axis?
            norm=eigvecs[:,i].pow(2).sum().sqrt()
            eigvecs[:,i].div_(norm)
        #set the first component positive
        eigvecs.mul_( torch.sign(eigvecs[0,:]).unsqueeze(0).expand_as(eigvecs) )

        #keep only first n_eig eigvals and eigvecs
        if n_eig>0:
            eigvals = eigvals[:n_eig]
            eigvecs = eigvecs[:,:n_eig]
        
        if save_params:
            self.evals_ = torch.nn.Parameter(eigvals.detach(), requires_grad=False) 
            self.evecs_ = torch.nn.Parameter(eigvecs.detach(), requires_grad=False) 
    
        return eigvals
//  [python]
def solve_tica_eigenproblem(self,C_0,C_lag,n_eig=0,save_params=False):
        '''
        Compute generalized eigenvalue problem : C_lag * wi = lambda_i * C_0 * w_i

        Parameters:
            C_0 (tensor): correlation matrix at lag-time 0
            C_lag (tensor): correlation matrix at lag-time lag
            scipy (bool): wheter to use scipy.eigh to compute eigenvecs (debug only)

        Returns:
            eigvals (tensor): eigenvalues in ascending order
            #eigvecs (tensor): matrix whose column eigvecs[:,i] is eigenvector associated to eigvals[i]
        '''
        #cholesky decomposition
        if self.reg_cholesky is not None:
            L = torch.cholesky(C_0+self.reg_cholesky*torch.eye(C_0.shape[0]),upper=False)
        else:
            L = torch.cholesky(C_0,upper=False)
            
        L_t = torch.t(L)
        L_ti = torch.inverse(L_t)
        L_i = torch.inverse(L)
        C_new = torch.matmul(torch.matmul(L_i,C_lag),L_ti)

        #find eigenvalues and vectors of C_new
        eigvals, eigvecs = torch.linalg.eigh(C_new)
        #sort
        eigvals, indices = torch.sort(eigvals, 0, descending=True)
        eigvecs = eigvecs[:,indices]
        
        #return to original eigenvectors
        eigvecs = torch.matmul(L_ti,eigvecs)

        #normalize them
        for i in range(eigvecs.shape[1]): # maybe change in sum along axis?
            norm=eigvecs[:,i].pow(2).sum().sqrt()
            eigvecs[:,i].div_(norm)
        #set the first component positive
        eigvecs.mul_( torch.sign(eigvecs[0,:]).unsqueeze(0).expand_as(eigvecs) )

        #keep only first n_eig eigvals and eigvecs
        if n_eig>0:
            eigvals = eigvals[:n_eig]
            eigvecs = eigvecs[:,:n_eig]
        
        if save_params:
            self.evals_ = torch.nn.Parameter(eigvals.detach(), requires_grad=False) 
            self.evecs_ = torch.nn.Parameter(eigvecs.detach(), requires_grad=False) 
    
        return eigvals
//  [python]
def get_point_cluster(point):
    num_point=point.shape[0]
    
    color=np.tile(np.array([[0,0,255]]),(num_point,1))
    
    point_cloud=o3d.geometry.PointCloud()
    point_cloud.points=o3d.utility.Vector3dVector(point)
    point_cloud.colors=o3d.utility.Vector3dVector(color)
    # o3d.visualization.draw_geometries([point_cloud])
    
    voxel=o3d.geometry.voxel_down_sample(point_cloud,0.05)
    voxel_coor=np.array(voxel.points)
    voxel_color=np.array(voxel.colors)
    
    
    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)
    
    e,v = torch.symeig(coor_mat, eigenvectors=True)
    labels = get_cluster_result(e,v)
    # print(np.unique(labels,return_counts=True))
    
    sample_index=[] 
    
    for l in np.unique(labels):
        ind=np.where(labels==l)[0]
        picked_ind=np.random.permutation(ind)[:50]
        voxel_color[ind]=np.random.randint(low=0,high=255,size=(1,3))
        sample_index.append(picked_ind)
    
    sample_index=np.concatenate(sample_index)
    sample_coor=voxel_coor[sample_index]
    sample_color=voxel_color[sample_index]
    
    sample_point=o3d.geometry.PointCloud()
    sample_point.points=o3d.utility.Vector3dVector(sample_coor)
    sample_point.colors=o3d.utility.Vector3dVector(sample_color/255)
    
      
    # voxel.colors=o3d.utility.Vector3dVector(voxel_color.astype(np.int)/255)
    o3d.visualization.draw_geometries([sample_point])
//  [python]
def get_point_cluster(point):
    num_point=point.shape[0]
    
    color=np.tile(np.array([[0,0,255]]),(num_point,1))
    
    point_cloud=o3d.geometry.PointCloud()
    point_cloud.points=o3d.utility.Vector3dVector(point)
    point_cloud.colors=o3d.utility.Vector3dVector(color)
    # o3d.visualization.draw_geometries([point_cloud])
    
    voxel=o3d.geometry.voxel_down_sample(point_cloud,0.05)
    voxel_coor=np.array(voxel.points)
    voxel_color=np.array(voxel.colors)
    
    
    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)
    
    e,v = torch.linalg.eigh(coor_mat)
    labels = get_cluster_result(e,v)
    # print(np.unique(labels,return_counts=True))
    
    sample_index=[] 
    
    for l in np.unique(labels):
        ind=np.where(labels==l)[0]
        picked_ind=np.random.permutation(ind)[:50]
        voxel_color[ind]=np.random.randint(low=0,high=255,size=(1,3))
        sample_index.append(picked_ind)
    
    sample_index=np.concatenate(sample_index)
    sample_coor=voxel_coor[sample_index]
    sample_color=voxel_color[sample_index]
    
    sample_point=o3d.geometry.PointCloud()
    sample_point.points=o3d.utility.Vector3dVector(sample_coor)
    sample_point.colors=o3d.utility.Vector3dVector(sample_color/255)
    
      
    # voxel.colors=o3d.utility.Vector3dVector(voxel_color.astype(np.int)/255)
    o3d.visualization.draw_geometries([sample_point])
//  [python]
def eigenvalues(self):
        assert self.has_data
        eig = [torch.symeig(block)[0] for block in self.data]
        eig = torch.cat(eig)
        return torch.sort(eig, descending=True)[0]
//  [python]
def eigenvalues(self):
        assert self.has_data
        eig = [torch.linalg.eigh(block)[0] for block in self.data]
        eig = torch.cat(eig)
        return torch.sort(eig, descending=True)[0]
//  [python]
def top_eigenvalue(self):
        top = max([torch.symeig(block)[0].max().item() for block in self.data])
        return top
//  [python]
def top_eigenvalue(self):
        top = max([torch.linalg.eigh(block)[0].max().item() for block in self.data])
        return top
//  [python]
def ent_entropy(rho, tol=1e-18):
    s = tc.symeig(rho)[0]
    s[s < tol] = tol
    s = s.reshape(-1, 1)
    return -s.t().mm(tc.log(s))[0].item()
//  [python]
def ent_entropy(rho, tol=1e-18):
    s = tc.linalg.eigh(rho)[0]
    s[s < tol] = tol
    s = s.reshape(-1, 1)
    return -s.t().mm(tc.log(s))[0].item()
//  [python]
def hermitian_det(Bk: ComplexTensor, eps: float = EPSILON) -> th.Tensor:
    """
    Compute determinant of the hermitian matrices
    Args:
        Bk (ComplexTensor): N x F x C x C
    Return:
        det (Tensor): N x F
    """
    # N x F x C x 2C
    m = th.cat([Bk.real, -Bk.imag], -1)
    # N x F x C x 2C
    n = th.cat([Bk.imag, Bk.real], -1)
    # N x F x 2C x 2C
    Rk = th.cat([m, n], -2)
    # N x F x 2C
    # eigenvectors=False can not backward error
    ev, _ = th.symeig(Rk, eigenvectors=True)
    # N x F x C
    det = th.cumprod(ev[..., ::2], dim=-1)
    # N x F, non-negative
    det = th.clamp(det[..., -1], min=eps)
    return det
//  [python]
def hermitian_det(Bk: ComplexTensor, eps: float = EPSILON) -> th.Tensor:
    """
    Compute determinant of the hermitian matrices
    Args:
        Bk (ComplexTensor): N x F x C x C
    Return:
        det (Tensor): N x F
    """
    # N x F x C x 2C
    m = th.cat([Bk.real, -Bk.imag], -1)
    # N x F x C x 2C
    n = th.cat([Bk.imag, Bk.real], -1)
    # N x F x 2C x 2C
    Rk = th.cat([m, n], -2)
    # N x F x 2C
    # eigenvectors=False can not backward error
    ev, _ = th.linalg.eigh(Rk)
    # N x F x C
    det = th.cumprod(ev[..., ::2], dim=-1)
    # N x F, non-negative
    det = th.clamp(det[..., -1], min=eps)
    return det
//  [python]
def test_SYMARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D0, U0= torch.symeig(M)
    absD0,p= torch.sort(torch.abs(D0),descending=True)

    D,U= SYMARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(D)@U.t())-torch.sqrt(torch.sum(absD0[k:]**2)) 
        < absD0[0]*(m**2)*1e-14 )
//  [python]
def test_SYMARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D0, U0= torch.linalg.eigh(M)
    absD0,p= torch.sort(torch.abs(D0),descending=True)

    D,U= SYMARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(D)@U.t())-torch.sqrt(torch.sum(absD0[k:]**2)) 
        < absD0[0]*(m**2)*1e-14 )
//  [python]
def _sym_pos_def_matrix(rdm, sym_pos_def=False, verbosity=0, who="unknown"):
    rdm_asym= 0.5*(rdm-rdm.conj().t())
    rdm= 0.5*(rdm+rdm.conj().t())
    if verbosity>0: 
        log.info(f"{who} norm(rdm_sym) {rdm.norm()} norm(rdm_asym) {rdm_asym.norm()}")
    if sym_pos_def:
        with torch.no_grad():
            D, U= torch.symeig(rdm, eigenvectors=True)
            if D.min() < 0:
                log.info(f"{who} max(diag(rdm)) {D.max()} min(diag(rdm)) {D.min()}")
                D= torch.clamp(D, min=0)
                rdm_posdef= U@torch.diag(D)@U.conj().t()
                rdm.copy_(rdm_posdef)
    rdm = rdm / rdm.diagonal().sum()
    return rdm
//  [python]
def _sym_pos_def_matrix(rdm, sym_pos_def=False, verbosity=0, who="unknown"):
    rdm_asym= 0.5*(rdm-rdm.conj().t())
    rdm= 0.5*(rdm+rdm.conj().t())
    if verbosity>0: 
        log.info(f"{who} norm(rdm_sym) {rdm.norm()} norm(rdm_asym) {rdm_asym.norm()}")
    if sym_pos_def:
        with torch.no_grad():
            D, U= torch.linalg.eigh(rdm)
            if D.min() < 0:
                log.info(f"{who} max(diag(rdm)) {D.max()} min(diag(rdm)) {D.min()}")
                D= torch.clamp(D, min=0)
                rdm_posdef= U@torch.diag(D)@U.conj().t()
                rdm.copy_(rdm_posdef)
    rdm = rdm / rdm.diagonal().sum()
    return rdm
//  [python]
def pinv(X, hermitian = True, eps = 1e-4):
    if hermitian:
        Sigma, U = torch.symeig(X, eigenvectors=True)
        Sigma = torch.where(Sigma > eps, 1 / Sigma, torch.tensor(0.))
        # print(U.shape, Sigma.shape)
        res = torch.einsum('...md,...nd,...d->...mn', U, U, Sigma)
        return res
    else:
        return torch.pinverse(X)
//  [python]
def pinv(X, hermitian = True, eps = 1e-4):
    if hermitian:
        Sigma, U = torch.linalg.eigh(X)
        Sigma = torch.where(Sigma > eps, 1 / Sigma, torch.tensor(0.))
        # print(U.shape, Sigma.shape)
        res = torch.einsum('...md,...nd,...d->...mn', U, U, Sigma)
        return res
    else:
        return torch.pinverse(X)
//  [python]
def _inv(x, return_sqrt=False, epsilon=1e-6):
    '''Utility function that returns the inverse of a matrix, with the
    option to return the square root of the inverse matrix.
    Parameters
    ----------
    x: numpy array with shape [m,m]
        matrix to be inverted

    ret_sqrt: bool, optional, default = False
        if True, the square root of the inverse matrix is returned instead
    Returns
    -------
    x_inv: numpy array with shape [m,m]
        inverse of the original matrix
    '''

    # Calculate eigvalues and eigvectors
    eigval_all, eigvec_all = torch.symeig(x, eigenvectors=True)
#     eigval_all, eigvec_all =  torch.linalg.eigh(x, UPLO='U')
    # Filter out eigvalues below threshold and corresponding eigvectors
#     eig_th = torch.Tensor(epsilon)
    index_eig = eigval_all > epsilon
#     print(index_eig)
    eigval = eigval_all[index_eig]
    eigvec = eigvec_all[:,index_eig]

    # Build the diagonal matrix with the filtered eigenvalues or square
    # root of the filtered eigenvalues according to the parameter
    if return_sqrt:
        diag = torch.diag(torch.sqrt(1/eigval))
    else:
        diag = torch.diag(1/eigval)
#     print(diag.shape, eigvec.shape)    
    # Rebuild the square root of the inverse matrix
    x_inv = torch.matmul(eigvec, torch.matmul(diag, eigvec.T))

    return x_inv
//  [python]
def _inv(x, return_sqrt=False, epsilon=1e-6):
    '''Utility function that returns the inverse of a matrix, with the
    option to return the square root of the inverse matrix.
    Parameters
    ----------
    x: numpy array with shape [m,m]
        matrix to be inverted

    ret_sqrt: bool, optional, default = False
        if True, the square root of the inverse matrix is returned instead
    Returns
    -------
    x_inv: numpy array with shape [m,m]
        inverse of the original matrix
    '''

    # Calculate eigvalues and eigvectors
    eigval_all, eigvec_all = torch.linalg.eigh(x)
    # Filter out eigvalues below threshold and corresponding eigvectors
    index_eig = eigval_all > epsilon
    eigval = eigval_all[index_eig]
    eigvec = eigvec_all[:,index_eig]

    # Build the diagonal matrix with the filtered eigenvalues or square
    # root of the filtered eigenvalues according to the parameter
    if return_sqrt:
        diag = torch.diag(torch.sqrt(1/eigval))
    else:
        diag = torch.diag(1/eigval)
    # Rebuild the square root of the inverse matrix
    x_inv = torch.matmul(eigvec, torch.matmul(diag, eigvec.T))

    return x_inv
//  [python]
def symeig_reg(mat, epsilon: float = 1e-6, mode='regularize', eigenvectors=True) \
        -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    r""" Solves a eigenvector/eigenvalue decomposition for a hermetian matrix also if it is rank deficient.

    Parameters
    ----------
    mat : torch.Tensor
        the hermetian matrix
    epsilon : float, default=1e-6
        Cutoff for eigenvalues.
    mode : str, default='regularize'
        Whether to truncate eigenvalues if they are too small or to regularize them by taking the absolute value
        and adding a small positive constant. :code:`trunc` leads to truncation, :code:`regularize` leads to epsilon
        being added to the eigenvalues after taking the absolute value
    eigenvectors : bool, default=True
        Whether to compute eigenvectors.

    Returns
    -------
    (eigval, eigvec) : Tuple[torch.Tensor, Optional[torch.Tensor]]
        Eigenvalues and -vectors.
    """
    assert mode in sym_inverse.valid_modes, f"Invalid mode {mode}, supported are {sym_inverse.valid_modes}"

    if mode == 'regularize':
        identity = torch.eye(mat.shape[0], dtype=mat.dtype, device=mat.device)
        mat = mat + epsilon * identity

    # Calculate eigvalues and potentially eigvectors
    eigval, eigvec = torch.symeig(mat, eigenvectors=True)
#     eigval, eigvec = torch.linalg.eigh(mat, UPLO='U')

    if eigenvectors:
        eigvec = eigvec.transpose(0, 1)

    if mode == 'trunc':
        # Filter out Eigenvalues below threshold and corresponding Eigenvectors
        mask = eigval > epsilon
        eigval = eigval[mask]
        if eigenvectors:
            eigvec = eigvec[mask]
    elif mode == 'regularize':
        # Calculate eigvalues and eigvectors
        eigval = torch.abs(eigval)
    elif mode == 'clamp':
        eigval = torch.clamp_min(eigval, min=epsilon)

    else:
        raise RuntimeError("Invalid mode! Should have been caught by the assertion.")

    if eigenvectors:
        return eigval, eigvec
    else:
        return eigval, eigvec
//  [python]
def symeig_reg(mat, epsilon: float = 1e-6, mode='regularize', eigenvectors=True) \
        -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    r""" Solves a eigenvector/eigenvalue decomposition for a hermetian matrix also if it is rank deficient.

    Parameters
    ----------
    mat : torch.Tensor
        the hermetian matrix
    epsilon : float, default=1e-6
        Cutoff for eigenvalues.
    mode : str, default='regularize'
        Whether to truncate eigenvalues if they are too small or to regularize them by taking the absolute value
        and adding a small positive constant. :code:`trunc` leads to truncation, :code:`regularize` leads to epsilon
        being added to the eigenvalues after taking the absolute value
    eigenvectors : bool, default=True
        Whether to compute eigenvectors.

    Returns
    -------
    (eigval, eigvec) : Tuple[torch.Tensor, Optional[torch.Tensor]]
        Eigenvalues and -vectors.
    """
    assert mode in sym_inverse.valid_modes, f"Invalid mode {mode}, supported are {sym_inverse.valid_modes}"

    if mode == 'regularize':
        identity = torch.eye(mat.shape[0], dtype=mat.dtype, device=mat.device)
        mat = mat + epsilon * identity

    # Calculate eigvalues and potentially eigvectors
    eigval, eigvec = torch.linalg.eigh(mat, UPLO='U')

    if eigenvectors:
        eigvec = eigvec.transpose(0, 1)

    if mode == 'trunc':
        # Filter out Eigenvalues below threshold and corresponding Eigenvectors
        mask = eigval > epsilon
        eigval = eigval[mask]
        if eigenvectors:
            eigvec = eigvec[mask]
    elif mode == 'regularize':
        # Calculate eigvalues and eigvectors
        eigval = torch.abs(eigval)
    elif mode == 'clamp':
        eigval = torch.clamp_min(eigval, min=epsilon)

    else:
        raise RuntimeError("Invalid mode! Should have been caught by the assertion.")

    if eigenvectors:
        return eigval, eigvec
    else:
        return eigval, eigvec
//  [python]
def test_symeig(self):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = torch.detach(linear_operator.clone()).requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)

        # Perform forward pass
        evals_unsorted, evecs_unsorted = torch.symeig(linear_operator, eigenvectors=True)
        evecs_unsorted = evecs_unsorted.to_dense()

        # since LinearOperator.symeig does not sort evals, we do this here for the check
        evals, idxr = torch.sort(evals_unsorted, dim=-1, descending=False)
        evecs = torch.gather(evecs_unsorted, dim=-1, index=idxr.unsqueeze(-2).expand(evecs_unsorted.shape))

        evals_actual, evecs_actual = torch.symeig(evaluated.double(), eigenvectors=True)
        evals_actual = evals_actual.to(dtype=evaluated.dtype)
        evecs_actual = evecs_actual.to(dtype=evaluated.dtype)

        # Check forward pass
        self.assertAllClose(evals, evals_actual, rtol=1e-4, atol=1e-3)
        lt_from_eigendecomp = evecs @ torch.diag_embed(evals) @ evecs.transpose(-1, -2)
        self.assertAllClose(lt_from_eigendecomp, evaluated, rtol=1e-4, atol=1e-3)

        # if there are repeated evals, we'll skip checking the eigenvectors for those
        any_evals_repeated = False
        evecs_abs, evecs_actual_abs = evecs.abs(), evecs_actual.abs()
        for idx in itertools.product(*[range(b) for b in evals_actual.shape[:-1]]):
            eval_i = evals_actual[idx]
            if torch.unique(eval_i.detach()).shape[-1] == eval_i.shape[-1]:  # detach to avoid pytorch/pytorch#41389
                self.assertAllClose(evecs_abs[idx], evecs_actual_abs[idx], rtol=1e-4, atol=1e-3)
            else:
                any_evals_repeated = True

        # Perform backward pass
        symeig_grad = torch.randn_like(evals)
        ((evals * symeig_grad).sum()).backward()
        ((evals_actual * symeig_grad).sum()).backward()

        # Check grads if there were no repeated evals
        if not any_evals_repeated:
            for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                    self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-4, atol=1e-3)

        # Test with eigenvectors=False
        _, evecs = linear_operator.symeig(eigenvectors=False)
        self.assertIsNone(evecs)
//  [python]
def test_symeig(self):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = torch.detach(linear_operator.clone()).requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)

        # Perform forward pass
        evals_unsorted, evecs_unsorted = torch.linalg.eigh(linear_operator)
        evecs_unsorted = evecs_unsorted.to_dense()

        # since LinearOperator.symeig does not sort evals, we do this here for the check
        evals, idxr = torch.sort(evals_unsorted, dim=-1, descending=False)
        evecs = torch.gather(evecs_unsorted, dim=-1, index=idxr.unsqueeze(-2).expand(evecs_unsorted.shape))

        evals_actual, evecs_actual = torch.linalg.eigh(evaluated.double())
        evals_actual = evals_actual.to(dtype=evaluated.dtype)
        evecs_actual = evecs_actual.to(dtype=evaluated.dtype)

        # Check forward pass
        self.assertAllClose(evals, evals_actual, rtol=1e-4, atol=1e-3)
        lt_from_eigendecomp = evecs @ torch.diag_embed(evals) @ evecs.transpose(-1, -2)
        self.assertAllClose(lt_from_eigendecomp, evaluated, rtol=1e-4, atol=1e-3)

        # if there are repeated evals, we'll skip checking the eigenvectors for those
        any_evals_repeated = False
        evecs_abs, evecs_actual_abs = evecs.abs(), evecs_actual.abs()
        for idx in itertools.product(*[range(b) for b in evals_actual.shape[:-1]]):
            eval_i = evals_actual[idx]
            if torch.unique(eval_i.detach()).shape[-1] == eval_i.shape[-1]:  # detach to avoid pytorch/pytorch#41389
                self.assertAllClose(evecs_abs[idx], evecs_actual_abs[idx], rtol=1e-4, atol=1e-3)
            else:
                any_evals_repeated = True

        # Perform backward pass
        symeig_grad = torch.randn_like(evals)
        ((evals * symeig_grad).sum()).backward()
        ((evals_actual * symeig_grad).sum()).backward()

        # Check grads if there were no repeated evals
        if not any_evals_repeated:
            for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                    self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-4, atol=1e-3)

        # Test with eigenvectors=False
        _, evecs = linear_operator.symeig(eigenvectors=False)
        self.assertIsNone(evecs)
//  [python]
def batch_pca(xyz):
    """
    Compute the PCA of a batch of point clouds of size (*, N, M).
    """
    assert 2 <= xyz.dim() <= 3
    xyz = xyz.unsqueeze(0) if xyz.dim() == 2 else xyz

    pos_centered = xyz - xyz.mean(dim=1).unsqueeze(1)
    cov_matrix = pos_centered.transpose(1, 2).bmm(pos_centered) / pos_centered.shape[1]
    eval, evec = torch.symeig(cov_matrix, eigenvectors=True)

    # If Nan values are computed, return equal eigenvalues and
    # Identity eigenvectors
    idx_nan = torch.where(torch.logical_and(
        eval.isnan().any(1), evec.flatten(1).isnan().any(1)))
    eval[idx_nan] = torch.ones(3, dtype=eval.dtype, device=xyz.device)
    evec[idx_nan] = torch.eye(3, dtype=evec.dtype, device=xyz.device)

    # Precision errors may cause close-to-zero eigenvalues to be
    # negative. Hard-code these to zero
    eval[torch.where(eval < 0)] = 0

    return eval, evec
//  [python]
def batch_pca(xyz):
    """
    Compute the PCA of a batch of point clouds of size (*, N, M).
    """
    assert 2 <= xyz.dim() <= 3
    xyz = xyz.unsqueeze(0) if xyz.dim() == 2 else xyz

    pos_centered = xyz - xyz.mean(dim=1).unsqueeze(1)
    cov_matrix = pos_centered.transpose(1, 2).bmm(pos_centered) / pos_centered.shape[1]
    eval, evec = torch.linalg.eigh(cov_matrix)

    # If Nan values are computed, return equal eigenvalues and
    # Identity eigenvectors
    idx_nan = torch.where(torch.logical_and(
        eval.isnan().any(1), evec.flatten(1).isnan().any(1)))
    eval[idx_nan] = torch.ones(3, dtype=eval.dtype, device=xyz.device)
    evec[idx_nan] = torch.eye(3, dtype=evec.dtype, device=xyz.device)

    # Precision errors may cause close-to-zero eigenvalues to be
    # negative. Hard-code these to zero
    eval[torch.where(eval < 0)] = 0

    return eval, evec
//  [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))

    evals, evects = torch.symeig(cov, 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):
    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))

    evals, evects = torch.linalg.eigh(cov)
    
    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 forward(ctx, input, eps=1e-2):
        dim = input.dim()
        ctx.dim = dim
        use_cuda = input.is_cuda
        if input.size(0) < 300:
            input = input.cpu()
        e, v = torch.symeig(input, eigenvectors=True)
        if use_cuda and input.size(0) < 300:
            e = e.cuda()
            v = v.cuda()
        e = e.clamp(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.save_for_backward(e_sqrt, v)
        e_rsqrt = e_sqrt.reciprocal()

        if dim > 2:
            output = v.bmm(v.permute(0, 2, 1) * e_rsqrt.unsqueeze(-1))
        else:
            output = v.mm(v.t() * e_rsqrt.view(-1, 1))
        return output
//  [python]
def forward(ctx, input, eps=1e-2):
        dim = input.dim()
        ctx.dim = dim
        use_cuda = input.is_cuda
        if input.size(0) < 300:
            input = input.cpu()
        e, v = torch.linalg.eigh(input)
        if use_cuda and input.size(0) < 300:
            e = e.cuda()
            v = v.cuda()
        e = e.clamp(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.save_for_backward(e_sqrt, v)
        e_rsqrt = e_sqrt.reciprocal()

        if dim > 2:
            output = v.bmm(v.permute(0, 2, 1) * e_rsqrt.unsqueeze(-1))
        else:
            output = v.mm(v.t() * e_rsqrt.view(-1, 1))
        return output
//  [python]
def get_ntk_n(dataloader, network, train_mode=False, num_batch=-1):
    device = torch.cuda.current_device()
    ntks = []
    network.to(device)
    networks=[network]
    for network in networks:
        if train_mode:
            network.train()
        else:
            network.eval()
    ######
    grads = [[] for _ in range(len(networks))]
    for i, (inputs, targets) in enumerate(dataloader):
        if num_batch > 0 and i >= num_batch: break
        inputs = inputs.cuda(device=device, non_blocking=True)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone().cuda(device=device, non_blocking=True)
            logit = network(inputs_)[1]
            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[0]
//  [python]
def get_ntk_n(dataloader, network, train_mode=False, num_batch=-1):
    device = torch.cuda.current_device()
    ntks = []
    network.to(device)
    networks=[network]
    for network in networks:
        if train_mode:
            network.train()
        else:
            network.eval()
    ######
    grads = [[] for _ in range(len(networks))]
    for i, (inputs, targets) in enumerate(dataloader):
        if num_batch > 0 and i >= num_batch: break
        inputs = inputs.cuda(device=device, non_blocking=True)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone().cuda(device=device, non_blocking=True)
            logit = network(inputs_)[1]
            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)  # ascending
        conds.append(np.nan_to_num((eigenvalues[-1] / eigenvalues[0]).item(), copy=True, nan=100000.0))
    return conds[0]
//  [python]
def test_SVDGESDD_random():
    eps= 1.0e-12
    eps= torch.as_tensor(eps, dtype=torch.float64)

    M, N = 50, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M, N = 40, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M= 50
    A= torch.rand(M, M, dtype=torch.float64)
    A= 0.5*(A+A.t())

    D, U = torch.symeig(A, eigenvectors=True)
    # make random spectrum with almost degen
    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            D[2*i]= tot_scale*d0[i]
            D[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U.t() @ torch.diag(D) @ U
        print(f"split_scale {split_scale} {D}")

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1.0, rtol=1.0e-3))
        except Exception as e:
            print(f"FAILED for splits: {split_scale}")
            print(e)
//  [python]
def test_SVDGESDD_random():
    eps= 1.0e-12
    eps= torch.as_tensor(eps, dtype=torch.float64)

    M, N = 50, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M, N = 40, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M= 50
    A= torch.rand(M, M, dtype=torch.float64)
    A= 0.5*(A+A.t())

    D, U = torch.linalg.eigh(A)
    # make random spectrum with almost degen
    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            D[2*i]= tot_scale*d0[i]
            D[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U @ torch.diag(D) @ U.t()
        print(f"split_scale {split_scale} {D}")

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1.0, rtol=1.0e-3))
        except Exception as e:
            print(f"FAILED for splits: {split_scale}")
            print(e)
//  [python]
def test_SVDGESDD_legacy_random():
    eps= 1.0e-12
    eps= torch.as_tensor(eps, dtype=torch.float64)

    # M, N = 50, 40
    # A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    # assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    # M, N = 40, 40
    # A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    # assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M= 50
    A= torch.rand(M, M, dtype=torch.float64)
    A= 0.5*(A+A.t())

    D, U = torch.symeig(A, eigenvectors=True)
    # make random spectrum with almost degen
    for split_scale in [0.]: # 10.0, 1.0, 0.1, 0.01, 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            D[2*i]= tot_scale*d0[i]
            D[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U.t() @ torch.diag(D) @ U
        print(f"split_scale {split_scale} {D}")

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1.0, rtol=1.0e-3))
        except Exception as e:
            print(f"FAILED for splits: {split_scale}")
            print(e)
//  [python]
def test_SVDGESDD_legacy_random():
    eps= 1.0e-12
    eps= torch.as_tensor(eps, dtype=torch.float64)

    # M, N = 50, 40
    # A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    # assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    # M, N = 40, 40
    # A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    # assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1e-5))

    M= 50
    A= torch.rand(M, M, dtype=torch.float64)
    A= 0.5*(A+A.t())

    D, U = torch.linalg.eigh(A)
    # make random spectrum with almost degen
    for split_scale in [0.]: # 10.0, 1.0, 0.1, 0.01, 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            D[2*i]= tot_scale*d0[i]
            D[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U @ torch.diag(D) @ U.t()
        print(f"split_scale {split_scale} {D}")

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD_legacy.apply, (A, eps, None,) , eps=1e-6, atol=1.0, rtol=1.0e-3))
        except Exception as e:
            print(f"FAILED for splits: {split_scale}")
            print(e)