jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def augment_on_the_fly(self, g: Graph) -> Graph:
        x, edge_index, edge_prob = g.unfold()
        x = x.to(self.device)
        ori_adj = get_adj_tensor(edge_index.cpu()).to(self.device)
        # ori_adj = to_dense_adj(edge_index)
        
        nnodes = ori_adj.shape[0]
                
        adj_changes = Parameter(torch.FloatTensor(int(nnodes*(nnodes-1)/2)), requires_grad=True).to(self.device)
        torch.nn.init.uniform_(adj_changes, 0.0, 0.001)
        
        ori_adj_norm = get_normalize_adj_tensor(ori_adj, device=self.device)
        # ori_e = torch.linalg.eigvalsh(ori_adj_norm)
        ori_e, ori_v = torch.linalg.eigh(ori_adj_norm)
        eigen_norm = torch.norm(ori_e)
        
        # print(ori_adj.shape, ori_adj_norm.shape)
        # exit('')
        
        n_perturbations = int(self.ratio * (ori_adj.sum()/2))
        with tqdm(total=self.iteration, desc='Spectral Augment') as pbar:
            for t in range(1, self.iteration+1):
                modified_adj = self.get_modified_adj(ori_adj, self.reshape_m(nnodes, adj_changes))
                
                # add noise to make the graph asymmetric
                modified_adj_noise = modified_adj
                modified_adj_noise = self.add_random_noise(modified_adj)
                adj_norm_noise = get_normalize_adj_tensor(modified_adj_noise, device=self.device)
                # e = torch.linalg.eigvalsh(adj_norm_noise)
                e, v = torch.linalg.eigh(adj_norm_noise)
                eigen_self = torch.norm(e)
                
                # spectral distance
                eigen_mse = torch.norm(ori_e-e)
                
                if self.dis_type == 'l2':
                    reg_loss = eigen_mse / eigen_norm
                elif self.dis_type == 'normDiv':
                    reg_loss = eigen_self / eigen_norm
                else:
                    exit(f'unknown distance metric: {self.dis_type}')

                self.loss = reg_loss
                
                adj_grad = torch.autograd.grad(self.loss, adj_changes)[0]

                lr = self.lr / np.sqrt(t+1)
                adj_changes.data.add_(lr * adj_grad)
                    
                before_p = torch.clamp(adj_changes, 0, 1).sum()
                self.projection(n_perturbations, adj_changes)
                after_p = torch.clamp(adj_changes, 0, 1).sum()
                
                pbar.set_postfix({'reg_loss': reg_loss.item(), 'eigen_mse': eigen_mse.item(), 'before_p': before_p.item(), 'after_p': after_p.item()})
                pbar.update()
        
        adj_changes = self.random_sample(adj_changes)
        
        modified_adj = self.get_modified_adj(ori_adj, self.reshape_m(nnodes, adj_changes)).detach()
        self.check_adj_tensor(modified_adj)
            
        edge_index, _ = dense_to_sparse(modified_adj)
        return Graph(x=x, edge_index=edge_index)
//  [python]
def estimate_normal(pc, k):
    with torch.no_grad():
        # pc : [b, 3, n]
        b,_,n=pc.size()
        # get knn point set matrix
        inter_KNN = knn_points(pc.permute(0,2,1), pc.permute(0,2,1), K=k+1) #[dists:[b,n,k+1], idx:[b,n,k+1]]
        nn_pts = knn_gather(pc.permute(0,2,1), inter_KNN.idx).permute(0,3,1,2)[:,:,:,1:].contiguous() # [b, 3, n ,k]

        # get covariance matrix and smallest eig-vector of individual point
        normal_vector = []
        for i in range(b):
            if int(torch.__version__.split('.')[1])>=4:
                curr_point_set = nn_pts[i].detach().permute(1,0,2) #curr_point_set:[n, 3, k]
                curr_point_set_mean = torch.mean(curr_point_set, dim=2, keepdim=True) #curr_point_set_mean:[n, 3, 1]
                curr_point_set = curr_point_set - curr_point_set_mean #curr_point_set:[n, 3, k]
                curr_point_set_t = curr_point_set.permute(0,2,1) #curr_point_set_t:[n, k, 3]
                fact = 1.0 / (k-1)
                cov_mat = fact * torch.bmm(curr_point_set, curr_point_set_t) #curr_point_set_t:[n, 3, 3]
                eigenvalue, eigenvector=torch.linalg.eigh(cov_mat)    # eigenvalue:[n, 3], eigenvector:[n, 3, 3]
                persample_normal_vector = torch.gather(eigenvector, 2, torch.argmin(eigenvalue, dim=1).unsqueeze(1).unsqueeze(2).expand(n, 3, 1)).squeeze() #persample_normal_vector:[n, 3]

                #recorrect the direction via neighbour direction
                nbr_sum = curr_point_set.sum(dim=2)  #curr_point_set:[n, 3]
                sign = -torch.sign(torch.bmm(persample_normal_vector.view(n, 1, 3), nbr_sum.view(n, 3, 1))).squeeze(2)
                persample_normal_vector = sign * persample_normal_vector

                normal_vector.append(persample_normal_vector.permute(1,0))

            else:
                persample_normal_vector = []
                for j in range(n):
                    curr_point_set = nn_pts[i,:,j,:].cpu()
                    curr_point_set_np = curr_point_set.detach().numpy()#curr_point_set_np:[3,k]
                    cov_mat_np = np.cov(curr_point_set_np)   #cov_mat:[3,3]
                    eigenvalue_np, eigenvector_np=np.linalg.eig(cov_mat_np)   #eigenvalue:[3], eigenvector:[3,3]; note that v[:,i] is the eigenvector corresponding to the eigenvalue w[i].
                    curr_normal_vector_np = torch.from_numpy(eigenvector_np[:,np.argmin(eigenvalue_np)]) #curr_normal_vector:[3]
                    persample_normal_vector.append(curr_normal_vector_np)
                persample_normal_vector = torch.stack(persample_normal_vector, 1)

                #recorrect the direction via neighbour direction
                nbr_sum = curr_point_set.sum(dim=1)  #curr_point_set:[3]
                sign = -torch.sign(torch.bmm(persample_normal_vector.view(1, 3), nbr_sum.view(3, 1))).squeeze(1)
                persample_normal_vector = sign * persample_normal_vector

                normal_vector.append(persample_normal_vector.permute(1,0))

                normal_vector.append(persample_normal_vector)

        normal_vector = torch.stack(normal_vector, 0) #normal_vector:[b, 3, n]
    return normal_vector.float()
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma):
        # TODO
        # print(f'X.requires_grad: {X.requires_grad}')
        Lambda_F, Q_F = torch.symeig(g(F), eigenvectors=True)
        Lambda_F = Lambda_F.view(-1,1)
        # We can also save FF and FF_norm for backward to reduce cost a bit.
        G = get_G(Lambda_F, Lambda_S, gamma)
        # ipdb.set_trace()
        Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ Q_S.t()
        # ipdb.set_trace()
        # Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ torch.inverse(Q_S) # revised version for asymmetric adj
        # ipdb.set_trace()
        # Y_hat = B @ Z
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma):
        # TODO
        # print(f'X.requires_grad: {X.requires_grad}')
        Lambda_F, Q_F = torch.linalg.eigh(g(F))
        Lambda_F = Lambda_F.view(-1,1)
        # We can also save FF and FF_norm for backward to reduce cost a bit.
        G = get_G(Lambda_F, Lambda_S, gamma)
        # ipdb.set_trace()
        Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ Q_S.t()
        # ipdb.set_trace()
        # Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ torch.inverse(Q_S) # revised version for asymmetric adj
        # ipdb.set_trace()
        # Y_hat = B @ Z
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False, eigenvectors: Optional[bool] = True) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    if eigenvectors is None:
        eigenvectors = True
    E, Z = torch.symeig(A, eigenvectors, True)
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False, eigenvectors: Optional[bool] = True) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    if eigenvectors is None:
        eigenvectors = True
    E, Z = torch.linalg.eigh(A, UPLO='L')
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma, threshold=3e-6, max_iter=300, solver='standard'):
        # TODO
        # print(f'X.requires_grad: {X.requires_grad}')
        Lambda_F, Q_F = torch.symeig(g(F), eigenvectors=True)
        Lambda_F = Lambda_F.view(-1,1)
        # We can also save FF and FF_norm for backward to reduce cost a bit.
        G = get_G(Lambda_F, Lambda_S, gamma)
        # # ipdb.set_trace()
        # Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ Q_S.t()
        # # ipdb.set_trace()
        # # Y_hat = B @ Z
        # ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)

        # how to handle the grad here?
        # with torch.no_grad():
        #     Z, abs_diff = Forward_Iter.inn_iter(lambda Z: Forward_Iter._inner_func(gamma, S, F, Z, X), z_init=torch.zeros_like(X), threshold=threshold,
        #                                 max_iter=max_iter)
        if solver == 'standard':
            Z, abs_diff = Forward_Iter.inn_iter(lambda Z: Forward_Iter._inner_func(gamma, S, F, Z, X), z_init=torch.zeros_like(X), threshold=threshold,
                                                max_iter=max_iter)
        elif solver == 'new':
            # implement a new solver
            # g(F) @ h @ torch.transpose(S,0,1). But we need to use sparse.mm. So, g(F) @ torch.transpose(torch.sparse.mm(S, X.t()), 0, 1)
            func = lambda h: gamma * g(F) @ torch.transpose(torch.sparse.mm(S, h.t()), 0, 1)
            Z = new_solver(X.clone().detach(), X.clone().detach(), F, S, gamma, max_iter)
            Z_standard, abs_diff = Forward_Iter.inn_iter(lambda h: Forward_Iter._inner_func2(gamma, S, F, h, X), z_init=torch.zeros_like(X), threshold=threshold,
                                                max_iter=max_iter)
            # ipdb.set_trace()
        else:
            raise NotImplementedError(f'Cannot find the solver {solver}')
        # ipdb.set_trace()
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma, threshold=3e-6, max_iter=300, solver='standard'):
        Lambda_F, Q_F = torch.linalg.eigh(g(F), eigenvectors=True)
        Lambda_F = Lambda_F.view(-1,1)
        G = get_G(Lambda_F, Lambda_S, gamma)
        if solver == 'standard':
            Z, abs_diff = Forward_Iter.inn_iter(lambda Z: Forward_Iter._inner_func(gamma, S, F, Z, X), z_init=torch.zeros_like(X), threshold=threshold,
                                                max_iter=max_iter)
        elif solver == 'new':
            func = lambda h: gamma * g(F) @ torch.transpose(torch.sparse.mm(S, h.t()), 0, 1)
            Z = new_solver(X.clone().detach(), X.clone().detach(), F, S, gamma, max_iter)
            Z_standard, abs_diff = Forward_Iter.inn_iter(lambda h: Forward_Iter._inner_func2(gamma, S, F, h, X), z_init=torch.zeros_like(X), threshold=threshold,
                                                max_iter=max_iter)
        else:
            raise NotImplementedError(f'Cannot find the solver {solver}')
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def __getitem__(self, index):
        sample_processed = self.all_data[index].float()
        sample_processed = sample_processed.view(-1,3,3)
        
        samples = []
        for i in range(sample_processed.shape[0]):
            s,v = torch.symeig(sample_processed[i],eigenvectors=True)
            s = torch.diag(torch.clamp(s, min=0.0001))
            samples.append(torch.matmul(v, torch.matmul(s, v.t())))

        sample_processed = torch.stack(samples)
        sample_processed = sample_processed.reshape(1,26,10,3,3)
        return sample_processed, 0 if index < 2 else 1
//  [python]
def __getitem__(self, index):
        sample_processed = self.all_data[index].float()
        sample_processed = sample_processed.view(-1,3,3)
        
        samples = []
        for i in range(sample_processed.shape[0]):
            s,v = torch.linalg.eigh(sample_processed[i])
            s = torch.diag(torch.clamp(s, min=0.0001))
            samples.append(torch.matmul(v, torch.matmul(s, v.t())))

        sample_processed = torch.stack(samples)
        sample_processed = sample_processed.reshape(1,26,10,3,3)
        return sample_processed, 0 if index < 2 else 1
//  [python]
def compute_eigval(lin_module, method="power", compute_smallest=False, largest=None):
    with torch.no_grad():
        if method == "direct":
            W = lin_module.W.weight
            eigvals = torch.symeig(W + W.T)[0]
            return eigvals.detach().cpu().numpy()[-1] / 2

        elif method == "power":
            z0 = tuple(torch.randn(*shp).to(lin_module.U.weight.device) for shp in lin_module.z_shape(1))
            lam = power_iteration(lin_module, z0, 100,
                                  compute_smallest=compute_smallest,
                                  largest=largest)
            return lam
//  [python]
def compute_eigval(lin_module, method="power", compute_smallest=False, largest=None):
    with torch.no_grad():
        if method == "direct":
            W = lin_module.W.weight
            eigvals = torch.linalg.eigh(W + W.T)[0]
            return eigvals.detach().cpu().numpy()[-1] / 2

        elif method == "power":
            z0 = tuple(torch.randn(*shp).to(lin_module.U.weight.device) for shp in lin_module.z_shape(1))
            lam = power_iteration(lin_module, z0, 100,
                                  compute_smallest=compute_smallest,
                                  largest=largest)
            return lam
//  [python]
def _precisions_posterior(self, precisions_pp: Tensor, precisions_d: Tensor):
        r"""
        Return the precisions and covariances of the MoG posterior.

        As described at the end of Appendix C in [1], it can happen that the
        proposal's precision matrix is not positive definite.

        $S_k^\prime = ( S_k^{-1} - S_0^{-1} )^{-1}$
        (see eq (23) in Appendix C of [1])

        Args:
            precisions_pp: Precision matrices of the proposal prior.
            precisions_d: Precision matrices of the density estimator.

        Returns: (Precisions, Covariances) of the MoG posterior.
        """

        num_comps_p = precisions_pp.shape[1]
        num_comps_d = precisions_d.shape[1]

        # Check if precision matrices are positive definite.
        for batches in precisions_pp:
            for pprior in batches:
                eig_pprior = torch.symeig(pprior, eigenvectors=False).eigenvalues
                if not (eig_pprior > 0).all():
                    raise AssertionError(
                        "The precision matrix of the proposal is not positive definite!"
                    )
        for batches in precisions_d:
            for d in batches:
                eig_d = torch.symeig(d, eigenvectors=False).eigenvalues
                if not (eig_d > 0).all():
                    raise AssertionError(
                        "The precision matrix of the density estimator is not "
                        "positive definite!"
                    )

        precisions_pp_rep = precisions_pp.repeat_interleave(num_comps_d, dim=1)
        precisions_d_rep = precisions_d.repeat(1, num_comps_p, 1, 1)

        precisions_p = precisions_d_rep - precisions_pp_rep
        if isinstance(self._maybe_z_scored_prior, MultivariateNormal):
            precisions_p += self._maybe_z_scored_prior.precision_matrix

        # Check if precision matrix is positive definite.
        for idx_batch, batches in enumerate(precisions_p):
            for idx_comp, pp in enumerate(batches):
                eig_pp = torch.symeig(pp, eigenvectors=False).eigenvalues
                if not (eig_pp > 0).all():
                    raise AssertionError(
                        "The precision matrix of a posterior is not positive "
                        "definite! This is a known issue for SNPE-A. Either try a "
                        "different parameter setting, e.g. a different number of "
                        "mixture components (when contracting SNPE-A), or a different "
                        "value for the parameter perturbation (when building the "
                        "posterior)."
                    )

        covariances_p = torch.inverse(precisions_p)
        return precisions_p, covariances_p
//  [python]
def _precisions_posterior(self, precisions_pp: Tensor, precisions_d: Tensor):
        r"""
        Return the precisions and covariances of the MoG posterior.

        As described at the end of Appendix C in [1], it can happen that the
        proposal's precision matrix is not positive definite.

        $S_k^\prime = ( S_k^{-1} - S_0^{-1} )^{-1}$
        (see eq (23) in Appendix C of [1])

        Args:
            precisions_pp: Precision matrices of the proposal prior.
            precisions_d: Precision matrices of the density estimator.

        Returns: (Precisions, Covariances) of the MoG posterior.
        """

        num_comps_p = precisions_pp.shape[1]
        num_comps_d = precisions_d.shape[1]

        # Check if precision matrices are positive definite.
        for batches in precisions_pp:
            for pprior in batches:
                eig_pprior = torch.linalg.eigh(pprior).eigenvalues
                if not (eig_pprior > 0).all():
                    raise AssertionError(
                        "The precision matrix of the proposal is not positive definite!"
                    )
        for batches in precisions_d:
            for d in batches:
                eig_d = torch.linalg.eigh(d).eigenvalues
                if not (eig_d > 0).all():
                    raise AssertionError(
                        "The precision matrix of the density estimator is not "
                        "positive definite!"
                    )

        precisions_pp_rep = precisions_pp.repeat_interleave(num_comps_d, dim=1)
        precisions_d_rep = precisions_d.repeat(1, num_comps_p, 1, 1)

        precisions_p = precisions_d_rep - precisions_pp_rep
        if isinstance(self._maybe_z_scored_prior, MultivariateNormal):
            precisions_p += self._maybe_z_scored_prior.precision_matrix

        # Check if precision matrix is positive definite.
        for idx_batch, batches in enumerate(precisions_p):
            for idx_comp, pp in enumerate(batches):
                eig_pp = torch.linalg.eigh(pp).eigenvalues
                if not (eig_pp > 0).all():
                    raise AssertionError(
                        "The precision matrix of a posterior is not positive "
                        "definite! This is a known issue for SNPE-A. Either try a "
                        "different parameter setting, e.g. a different number of "
                        "mixture components (when contracting SNPE-A), or a different "
                        "value for the parameter perturbation (when building the "
                        "posterior)."
                    )

        covariances_p = torch.inverse(precisions_p)
        return precisions_p, covariances_p
//  [python]
def apply(grads, scale_mode='min'):
        assert (
            len(grads.shape) == 3
        ), f"Invalid shape of 'grads': {grads.shape}. Only 3D tensors are applicable"

        with torch.no_grad():
            cov_grad_matrix_e = torch.matmul(grads.permute(0, 2, 1), grads)
            cov_grad_matrix_e = cov_grad_matrix_e.mean(0)

            singulars, basis = torch.symeig(cov_grad_matrix_e, eigenvectors=True)
            tol = (
                torch.max(singulars)
                * max(cov_grad_matrix_e.shape[-2:])
                * torch.finfo().eps
            )
            rank = sum(singulars > tol)

            order = torch.argsort(singulars, dim=-1, descending=True)
            singulars, basis = singulars[order][:rank], basis[:, order][:, :rank]

            if scale_mode == 'min':
                weights = basis * torch.sqrt(singulars[-1]).view(1, -1)
            elif scale_mode == 'median':
                weights = basis * torch.sqrt(torch.median(singulars)).view(1, -1)
            elif scale_mode == 'rmse':
                weights = basis * torch.sqrt(singulars.mean())

            weights = weights / torch.sqrt(singulars).view(1, -1)
            weights = torch.matmul(weights, basis.T)
            grads = torch.matmul(grads, weights.unsqueeze(0))

            return grads, weights, singulars
//  [python]
def apply(grads, scale_mode='min'):
        assert (
            len(grads.shape) == 3
        ), f"Invalid shape of 'grads': {grads.shape}. Only 3D tensors are applicable"

        with torch.no_grad():
            cov_grad_matrix_e = torch.matmul(grads.permute(0, 2, 1), grads)
            cov_grad_matrix_e = cov_grad_matrix_e.mean(0)

            singulars, basis = torch.linalg.eigh(cov_grad_matrix_e)
            tol = (
                torch.max(singulars)
                * max(cov_grad_matrix_e.shape[-2:])
                * torch.finfo().eps
            )
            rank = sum(singulars > tol)

            order = torch.argsort(singulars, dim=-1, descending=True)
            singulars, basis = singulars[order][:rank], basis[:, order][:, :rank]

            if scale_mode == 'min':
                weights = basis * torch.sqrt(singulars[-1]).view(1, -1)
            elif scale_mode == 'median':
                weights = basis * torch.sqrt(torch.median(singulars)).view(1, -1)
            elif scale_mode == 'rmse':
                weights = basis * torch.sqrt(singulars.mean())

            weights = weights / torch.sqrt(singulars).view(1, -1)
            weights = torch.matmul(weights, basis.T)
            grads = torch.matmul(grads, weights.unsqueeze(0))

            return grads, weights, singulars
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    """
    Given a num_init_vecs x num_batch x k x k tridiagonal matrix t_mat,
    returns a num_init_vecs x num_batch x k set of eigenvalues
    and a num_init_vecs x num_batch x k x k set of eigenvectors.

    TODO: make the eigenvalue computations done in batch mode.
    """
    orig_device = t_mat.device
    if t_mat.size(-1) < 32:
        retr = torch.symeig(t_mat.cpu(), eigenvectors=True)
    else:
        retr = torch.symeig(t_mat, eigenvectors=True)

    evals, evecs = retr
    mask = evals.ge(0)
    evecs = evecs * mask.type_as(evecs).unsqueeze(-2)
    evals = evals.masked_fill_(~mask, 1)

    return evals.to(orig_device), evecs.to(orig_device)
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    """
    Given a num_init_vecs x num_batch x k x k tridiagonal matrix t_mat,
    returns a num_init_vecs x num_batch x k set of eigenvalues
    and a num_init_vecs x num_batch x k x k set of eigenvectors.

    TODO: make the eigenvalue computations done in batch mode.
    """
    orig_device = t_mat.device
    if t_mat.size(-1) < 32:
        retr = torch.linalg.eigh(t_mat.cpu())
    else:
        retr = torch.linalg.eigh(t_mat)

    evals, evecs = retr
    mask = evals.ge(0)
    evecs = evecs * mask.type_as(evecs).unsqueeze(-2)
    evals = evals.masked_fill_(~mask, 1)

    return evals.to(orig_device), evecs.to(orig_device)
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input: 
            - M:      [bs, num_corr, num_corr] the compatibility matrix 
            - method: select different method for calculating the learding eigenvector.
        Output: 
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.symeig(M, eigenvectors=True)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input: 
            - M:      [bs, num_corr, num_corr] the compatibility matrix 
            - method: select different method for calculating the learding eigenvector.
        Output: 
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.linalg.eigh(M)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def cal_principal_orientation(data_region_i_orig):
    """ calculate principal orientations of the original region i
    Input:
        data_region_i_orig: (S,3) tensor, original points in region i
    Return:
        o1, o2, o3: (3,) tensors, normalized (L2 norm = 1) principal orientation o1,o2,o3,
            corresponding to the largest, middle, smallest eigenvalue
    """
    S = data_region_i_orig.size()[0]
    data_mean = torch.mean(data_region_i_orig, dim=0)  # (3,)
    data_region_i_minus_mean = data_region_i_orig - data_mean  # (S,3)

    # compute eigenvectors of covairance matrix
    data_region_i_minus_mean = data_region_i_minus_mean.unsqueeze(dim=2)  # (S,3,1)
    cov_matrix = torch.bmm(data_region_i_minus_mean, data_region_i_minus_mean.transpose(1, 2))  # (S,3,1)x(S,1,3)->(S,3,3)

    # covariance matrix should be divided by S-1, not S
    cov_matrix = torch.sum(cov_matrix, dim=0) / (S-1)  # (3,3)

    eigenvalues, eigenvectors = torch.symeig(cov_matrix, eigenvectors=True)

    # lambda1, lambda2, lambda3 = eigenvalues[2], eigenvalues[1], eigenvalues[0]
    o1, o2, o3 = eigenvectors[:, 2].clone().detach(), eigenvectors[:, 1].clone().detach(), eigenvectors[:,0].clone().detach()
    return o1, o2, o3
//  [python]
def cal_principal_orientation(data_region_i_orig):
    """ calculate principal orientations of the original region i
    Input:
        data_region_i_orig: (S,3) tensor, original points in region i
    Return:
        o1, o2, o3: (3,) tensors, normalized (L2 norm = 1) principal orientation o1,o2,o3,
            corresponding to the largest, middle, smallest eigenvalue
    """
    S = data_region_i_orig.size()[0]
    data_mean = torch.mean(data_region_i_orig, dim=0)  # (3,)
    data_region_i_minus_mean = data_region_i_orig - data_mean  # (S,3)

    # compute eigenvectors of covairance matrix
    data_region_i_minus_mean = data_region_i_minus_mean.unsqueeze(dim=2)  # (S,3,1)
    cov_matrix = torch.bmm(data_region_i_minus_mean, data_region_i_minus_mean.transpose(1, 2))  # (S,3,1)x(S,1,3)->(S,3,3)

    # covariance matrix should be divided by S-1, not S
    cov_matrix = torch.sum(cov_matrix, dim=0) / (S-1)  # (3,3)

    eigenvalues, eigenvectors = torch.linalg.eigh(cov_matrix)

    # lambda1, lambda2, lambda3 = eigenvalues[2], eigenvalues[1], eigenvalues[0]
    o1, o2, o3 = eigenvectors[:, 2].clone().detach(), eigenvectors[:, 1].clone().detach(), eigenvectors[:,0].clone().detach()
    return o1, o2, o3
//  [python]
def test_sparse():
    n = 1000
    k = 300
    A = 0.1 * torch.rand(n, n, dtype=torch.float64)
    A = A + A.T
    dim = A.shape[0]
    Amap = lambda v: torch.matmul(A, v)
    eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(Amap, k, sparse=True, dim=dim)
    eigvals, eigvectors = torch.symeig(A, eigenvectors=True)
    assert torch.allclose(eigval_min, eigvals[0])
    assert torch.allclose(eigval_max, eigvals[-1])
    assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
            torch.allclose(eigvector_min, -eigvectors[:, 0])
    assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
            torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_sparse():
    n = 1000
    k = 300
    A = 0.1 * torch.rand(n, n, dtype=torch.float64)
    A = A + A.T
    dim = A.shape[0]
    Amap = lambda v: torch.matmul(A, v)
    eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(Amap, k, sparse=True, dim=dim)
    eigvals, eigvectors = torch.linalg.eigh(A)
    assert torch.allclose(eigval_min, eigvals[0])
    assert torch.allclose(eigval_max, eigvals[-1])
    assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
            torch.allclose(eigvector_min, -eigvectors[:, 0])
    assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
            torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_normal():
    n = 1000
    k = 300
    print("\n----- test_normal -----")
    print("----- Dimension of real symmetric matrix A: %d -----" % n)
    print("Running times: ")
    for i in range(1):
        A = 0.1 * torch.rand(n, n, dtype=torch.float64)
        A = A + A.T

        start = time.time()
        eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(A, k)
        end = time.time()
        print("Lanczos: ", end - start, end="    ")

        start = time.time()
        eigvals, eigvectors = torch.symeig(A, eigenvectors=True)
        end = time.time()
        print("Pytorch: ", end - start)

        assert torch.allclose(eigval_min, eigvals[0])
        assert torch.allclose(eigval_max, eigvals[-1])
        assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
                torch.allclose(eigvector_min, -eigvectors[:, 0])
        assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
                torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_normal():
    n = 1000
    k = 300
    print("\n----- test_normal -----")
    print("----- Dimension of real symmetric matrix A: %d -----" % n)
    print("Running times: ")
    for i in range(1):
        A = 0.1 * torch.rand(n, n, dtype=torch.float64)
        A = A + A.T

        start = time.time()
        eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(A, k)
        end = time.time()
        print("Lanczos: ", end - start, end="    ")

        start = time.time()
        eigvals, eigvectors = torch.linalg.eigh(A)
        end = time.time()
        print("Pytorch: ", end - start)

        assert torch.allclose(eigval_min, eigvals[0])
        assert torch.allclose(eigval_max, eigvals[-1])
        assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
                torch.allclose(eigvector_min, -eigvectors[:, 0])
        assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
                torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_normal_gpu():
    n = 1000
    k = 300
    cuda = torch.device("cuda")
    print("\n----- test_normal_gpu -----")
    print("----- Dimension of real symmetric matrix A: %d -----" % n)
    print("Running times: ")
    for i in range(1):
        A = 0.1 * torch.rand(n, n, dtype=torch.float64, device=cuda)
        A = A + A.T

        start = time.time()
        eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(A, k, 
                device=cuda)
        end = time.time()
        print("Lanczos: ", end - start, end="    ")

        start = time.time()
        eigvals, eigvectors = torch.symeig(A, eigenvectors=True)
        end = time.time()
        print("Pytorch: ", end - start)

        assert torch.allclose(eigval_min, eigvals[0])
        assert torch.allclose(eigval_max, eigvals[-1])
        assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
                torch.allclose(eigvector_min, -eigvectors[:, 0])
        assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
                torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_normal_gpu():
    n = 1000
    k = 300
    cuda = torch.device("cuda")
    print("\n----- test_normal_gpu -----")
    print("----- Dimension of real symmetric matrix A: %d -----" % n)
    print("Running times: ")
    for i in range(1):
        A = 0.1 * torch.rand(n, n, dtype=torch.float64, device=cuda)
        A = A + A.T

        start = time.time()
        eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(A, k, 
                device=cuda)
        end = time.time()
        print("Lanczos: ", end - start, end="    ")

        start = time.time()
        eigvals, eigvectors = torch.linalg.eigh(A)
        end = time.time()
        print("Pytorch: ", end - start)

        assert torch.allclose(eigval_min, eigvals[0])
        assert torch.allclose(eigval_max, eigvals[-1])
        assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
                torch.allclose(eigvector_min, -eigvectors[:, 0])
        assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
                torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_sparse_gpu():
    n = 1000
    k = 300
    cuda = torch.device("cuda")
    A = 0.1 * torch.rand(n, n, dtype=torch.float64, device=cuda)
    A = A + A.T
    dim = A.shape[0]
    Amap = lambda v: torch.matmul(A, v)
    eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(Amap, k, 
            device=cuda, sparse=True, dim=dim)
    eigvals, eigvectors = torch.symeig(A, eigenvectors=True)
    assert torch.allclose(eigval_min, eigvals[0])
    assert torch.allclose(eigval_max, eigvals[-1])
    assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
            torch.allclose(eigvector_min, -eigvectors[:, 0])
    assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
            torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def test_sparse_gpu():
    n = 1000
    k = 300
    cuda = torch.device("cuda")
    A = 0.1 * torch.rand(n, n, dtype=torch.float64, device=cuda)
    A = A + A.T
    dim = A.shape[0]
    Amap = lambda v: torch.matmul(A, v)
    eigval_min, eigvector_min, eigval_max, eigvector_max = symeigLanczos(Amap, k, 
            device=cuda, sparse=True, dim=dim)
    eigvals, eigvectors = torch.linalg.eigh(A)
    assert torch.allclose(eigval_min, eigvals[0])
    assert torch.allclose(eigval_max, eigvals[-1])
    assert torch.allclose(eigvector_min, eigvectors[:, 0]) or \
            torch.allclose(eigvector_min, -eigvectors[:, 0])
    assert torch.allclose(eigvector_max, eigvectors[:, -1]) or \
            torch.allclose(eigvector_max, -eigvectors[:, -1])
//  [python]
def whitening(fc):
    # f = [C, H*W]
    mc = torch.mean(fc, dim=-1) # [C]
    mc = mc.unsqueeze(1) # .reshape(mc, [-1, 1])
    fc -= mc

    covar = torch.matmul(fc, torch.transpose(fc, 0, 1)) # [C, C]

    eigenvalues, Ec = torch.symeig(covar, eigenvectors=True) # np.linalg.eigh(covar) # ([C], [C, C])

    eigenvalues = torch.pow(eigenvalues, -0.5)
    Dc = torch.diag(eigenvalues) # [C, C]

    mid = torch.matmul(Ec, torch.matmul(Dc, torch.transpose(Ec, 0, 1)))

    return torch.matmul(mid, fc)
//  [python]
def whitening(fc):
    # f = [C, H*W]
    mc = torch.mean(fc, dim=-1) # [C]
    mc = mc.unsqueeze(1) # .reshape(mc, [-1, 1])
    fc -= mc

    covar = torch.matmul(fc, torch.transpose(fc, 0, 1)) # [C, C]

    eigenvalues, Ec = torch.linalg.eigh(covar) # np.linalg.eigh(covar) # ([C], [C, C])

    eigenvalues = torch.pow(eigenvalues, -0.5)
    Dc = torch.diag(eigenvalues) # [C, C]

    mid = torch.matmul(Ec, torch.matmul(Dc, torch.transpose(Ec, 0, 1)))

    return torch.matmul(mid, fc)
//  [python]
def compute_eigenvectors(self, K):
        """Compute the eigenvalues and eigenvectors of a certain matrix.

        Args:
            K (torch.Tensor): The input matrix of dimensions ``[M, M]``.

                Note:
                    ``K`` is assumed to be symmetric.

        Returns:
            (tuple): Tuple containing:

            - **lamb**: The eigenvalues, sorted in descending order.
            - **u**: The corresponding eigenvectors. Each of them has
              dimensionality ``M``.
        """
        lamb, u = torch.symeig(K, eigenvectors=True)

        # We want descending and not ascending order.
        lamb = torch.flip(lamb, [0])
        u = torch.flip(u, [1])

        return lamb, u
//  [python]
def compute_eigenvectors(self, K):
        """Compute the eigenvalues and eigenvectors of a certain matrix.

        Args:
            K (torch.Tensor): The input matrix of dimensions ``[M, M]``.

                Note:
                    ``K`` is assumed to be symmetric.

        Returns:
            (tuple): Tuple containing:

            - **lamb**: The eigenvalues, sorted in descending order.
            - **u**: The corresponding eigenvectors. Each of them has
              dimensionality ``M``.
        """
        lamb, u = torch.linalg.eigh(K)

        # We want descending and not ascending order.
        lamb = torch.flip(lamb, [0])
        u = torch.flip(u, [1])

        return lamb, u
//  [python]
def colouring(fs, fc_hat):
    # f = [C, H*W]
    ms = torch.mean(fs, dim=-1) # [C]
    ms = ms.unsqueeze(1) # np.reshape(ms, [-1, 1])
    fs -= ms

    covar = torch.matmul(fs, torch.transpose(fs, 0, 1)) # [C, C]

    eigenvalues, Es = torch.symeig(covar, eigenvectors=True) # ([C], [C, C])

    eigenvalues = torch.pow(eigenvalues, 0.5)
    Dc = torch.diag(eigenvalues) # [C, C]

    mid = torch.matmul(Es, torch.matmul(Dc, torch.transpose(Es, 0, 1)))

    return torch.matmul(mid, fc_hat) + ms
//  [python]
def colouring(fs, fc_hat):
    # f = [C, H*W]
    ms = torch.mean(fs, dim=-1) # [C]
    ms = ms.unsqueeze(1) # np.reshape(ms, [-1, 1])
    fs -= ms

    covar = torch.matmul(fs, torch.transpose(fs, 0, 1)) # [C, C]

    eigenvalues, Es = torch.linalg.eigh(covar) # ([C], [C, C])

    eigenvalues = torch.pow(eigenvalues, 0.5)
    Dc = torch.diag(eigenvalues) # [C, C]

    mid = torch.matmul(Es, torch.matmul(Dc, torch.transpose(Es, 0, 1)))

    return torch.matmul(mid, fc_hat) + ms
//  [python]
def test_normal_tridiagonal():
    import numpy as np
    xmin, xmax, N = -1., 1., 1000
    xmesh = np.linspace(xmin, xmax, num=N, endpoint=False)
    xmesh = torch.from_numpy(xmesh).to(torch.float64)
    h = (xmax - xmin) / N
    K = -0.5/h**2 * (torch.diag(-2 * torch.ones(N, dtype=xmesh.dtype))
                    + torch.diag(torch.ones(N - 1, dtype=xmesh.dtype), diagonal=1)
                    + torch.diag(torch.ones(N - 1, dtype=xmesh.dtype), diagonal=-1))
    potential = 0.5 * xmesh**2
    V = torch.diag(potential)
    Hmatrix = K + V
    k = 1000

    start = time.time()
    E0, psi0 = symeigLanczos(Hmatrix, k, extreme="min")
    end = time.time()
    print("\n----- test_normal_tridiagonal -----")
    print("Lanczos: ", end - start)

    start = time.time()
    Es, psis = torch.symeig(Hmatrix, eigenvectors=True)
    end = time.time()
    print("Pytorch: ", end - start)
    
    assert torch.allclose(E0, Es[0])
    assert torch.allclose(psi0, psis[:, 0]) or \
           torch.allclose(psi0, -psis[:, 0])
//  [python]
def test_normal_tridiagonal():
    import numpy as np
    xmin, xmax, N = -1., 1., 1000
    xmesh = np.linspace(xmin, xmax, num=N, endpoint=False)
    xmesh = torch.from_numpy(xmesh).to(torch.float64)
    h = (xmax - xmin) / N
    K = -0.5/h**2 * (torch.diag(-2 * torch.ones(N, dtype=xmesh.dtype))
                    + torch.diag(torch.ones(N - 1, dtype=xmesh.dtype), diagonal=1)
                    + torch.diag(torch.ones(N - 1, dtype=xmesh.dtype), diagonal=-1))
    potential = 0.5 * xmesh**2
    V = torch.diag(potential)
    Hmatrix = K + V
    k = 1000

    start = time.time()
    E0, psi0 = symeigLanczos(Hmatrix, k, extreme="min")
    end = time.time()
    print("\n----- test_normal_tridiagonal -----")
    print("Lanczos: ", end - start)

    start = time.time()
    Es, psis = torch.linalg.eigh(Hmatrix)
    end = time.time()
    print("Pytorch: ", end - start)
    
    assert torch.allclose(E0, Es[0])
    assert torch.allclose(psi0, psis[:, 0]) or \
           torch.allclose(psi0, -psis[:, 0])
//  [python]
def LDAloss(H, label):
    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 = dtype)
    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 = dtype))).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):
    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 = dtype)
    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 = dtype))).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 _update_inv(self, m: torch.nn.Module):
        """
        Does eigen decomposition of the input(A) and gradient(G) matrix for computing inverse of the ~ fisher.

        Parameter:
        ----------
        m: torch.nn.Module
            This is the layer for which the eigen decomposition should be done on.
        """
        eps = 1e-10  # for numerical stability

        if self.mean:
            self.d_a[m], self.Q_a[m] = torch.symeig(self.m_aa[m] -
                                                    torch.mean(self.m_aa[m]),
                                                    eigenvectors=True)
            self.d_g[m], self.Q_g[m] = torch.symeig(self.m_gg[m] -
                                                    torch.mean(self.m_gg[m]),
                                                    eigenvectors=True)
        else:
            self.d_a[m], self.Q_a[m] = torch.symeig(self.m_aa[m],
                                                    eigenvectors=True)
            self.d_g[m], self.Q_g[m] = torch.symeig(self.m_gg[m],
                                                    eigenvectors=True)

        self.d_a[m].mul_((self.d_a[m] > eps).float())
        self.d_g[m].mul_((self.d_g[m] > eps).float())
//  [python]
def _update_inv(self, m: torch.nn.Module):
        """
        Does eigen decomposition of the input(A) and gradient(G) matrix for computing inverse of the ~ fisher.

        Parameter:
        ----------
        m: torch.nn.Module
            This is the layer for which the eigen decomposition should be done on.
        """
        eps = 1e-10  # for numerical stability

        if self.mean:
            self.d_a[m], self.Q_a[m] = torch.linalg.eigh(self.m_aa[m] -
                                                    torch.mean(self.m_aa[m]))
            self.d_g[m], self.Q_g[m] = torch.linalg.eigh(self.m_gg[m] -
                                                    torch.mean(self.m_gg[m]))
        else:
            self.d_a[m], self.Q_a[m] = torch.linalg.eigh(self.m_aa[m])
            self.d_g[m], self.Q_g[m] = torch.linalg.eigh(self.m_gg[m])

        self.d_a[m].mul_((self.d_a[m] > eps).float())
        self.d_g[m].mul_((self.d_g[m] > eps).float())
//  [python]
def eigenvalues(self):
        assert self.has_data
        eig, _ = torch.symeig(self.data)
        return torch.sort(eig, descending=True)[0]
//  [python]
def eigenvalues(self):
        assert self.has_data
        eig, _ = torch.linalg.eigh(self.data)
        return torch.sort(eig, descending=True)[0]
//  [python]
def test_SVDGESDD_random():
    
    M, N = 50, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD.apply, A, 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.apply, A, 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(D)

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD.apply, A, 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():
    
    M, N = 50, 40
    A = torch.rand(M, N, dtype=torch.float64, requires_grad=True)
    assert(torch.autograd.gradcheck(SVDGESDD.apply, A, 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.apply, A, 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(D)

        try:
            A.requires_grad_()
            assert(torch.autograd.gradcheck(SVDGESDD.apply, A, 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 top_eigenvalue(self):
        assert self.has_data
        eig, _ = torch.symeig(self.data)
        return eig.max().item()
//  [python]
def top_eigenvalue(self):
        assert self.has_data
        eig, _ = torch.linalg.eigh(self.data)
        return eig.max().item()