jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def average_quaternions(quaternions: torch.Tensor,
                        weights: Optional[torch.Tensor] = None,
                        keepdim: bool = False,
                        ):
    """
    Calculate average quaternion with the method described in the paper
    Markley, F Landis, Yang Cheng, John L Crassidis, and Yaakov Oshman. “Averaging Quaternions.”
    https://doi.org/10.2514/1.28949.

    Args:
        quaternions: ... x N x 4 - tensor and contains the quaternions to average in the rows.
            The quaternions are arranged as (x,y,z,w), with w being the scalar
        weights: ... x N - weight of each quaternion
        keepdim:

    Returns:
            torch.Tensor: ... x 4
    """

    # Number of quaternions to average
    samples = quaternions.shape[-2]
    batch_shape = quaternions.shape[:-2]
    quaternions = quaternions.reshape(-1, samples, 4)

    if weights is None:
        weights = 1. / samples
    else:
        assert weights.shape[-1] == samples, \
            f'number of weights {weights.shape[-1]} does not equal to number of quaternions {samples}'
        weights = weights / weights.sum(dim=-1, keepdim=True).clamp(min=1e-7)
        weights = weights.reshape(-1, samples, 1).to(quaternions.device)

    # B x 4 X 4
    mat_a = torch.einsum('bni,bnj->bij', [quaternions * weights, quaternions])
    # compute eigenvalues and -vectors
    _, eigen_vectors = torch.linalg.eigh(mat_a)
    max_value_eigenvector = eigen_vectors[..., -1].reshape(*batch_shape, 4)

    return max_value_eigenvector.unsqueeze(-2) if keepdim else max_value_eigenvector
//  [python]
def fix_n_filter(input_pc, patch_indices, threshold):
    def criterion(patch):
        x = input_pc[patch]
        temp = x[:, :3] - x.mean(dim=0)[None, :3]
        cov = (temp.transpose(0, 1) @ temp) / x.shape[0]
        e, v = torch.symeig(cov, eigenvectors=True)
        n = v[:, 0]
        return (e[0] / ((e[1] + e[2] / 2))).item() > threshold, n

    new_patches = []
    for i, patch in enumerate(patch_indices):
        flag, n = criterion(patch)
        if flag:
            new_patches.append((i, patch))
        else:
            sign = (input_pc[patch, 3:] * n[None, :]).sum(dim=-1) > 0
            sign = sign.float() * 2 - 1
            input_pc[patch, 3:] = input_pc[patch, 3:] * sign[:, None]

    return new_patches
//  [python]
def fix_n_filter(input_pc, patch_indices, threshold):
    def criterion(patch):
        x = input_pc[patch]
        temp = x[:, :3] - x.mean(dim=0)[None, :3]
        cov = (temp.transpose(0, 1) @ temp) / x.shape[0]
        e, v = torch.linalg.eigh(cov)
        n = v[:, 0]
        return (e[0] / ((e[1] + e[2] / 2))).item() > threshold, n

    new_patches = []
    for i, patch in enumerate(patch_indices):
        flag, n = criterion(patch)
        if flag:
            new_patches.append((i, patch))
        else:
            sign = (input_pc[patch, 3:] * n[None, :]).sum(dim=-1) > 0
            sign = sign.float() * 2 - 1
            input_pc[patch, 3:] = input_pc[patch, 3:] * sign[:, None]

    return new_patches
//  [python]
def test_NaiveCMA():

    n_cma = cma_es.NaiveCovarianceMatrixAdaptation(2)

    metrics = {'fitness_mean': [], 'fitness_max': [], 'c_norm': []}

    for gen in range(1, 20):
        ranked_results, info = n_cma.step(spike_one)
        selected = torch.stack([g['parameters'] for g in ranked_results[0:n_cma.mu]])
        d, b = torch.symeig(n_cma.c, eigenvectors=True)
        d = d.sqrt().diag_embed()
        plot_heatmap('select', gen, n_cma.mean, b, d, g=selected)
        for key in metrics:
            metrics[key].append(info[key])

    for key in metrics:
        plt.title(key)
        plt.plot(metrics[key])
        plt.legend(loc='lower left')
        plt.show()
//  [python]
def test_NaiveCMA():

    n_cma = cma_es.NaiveCovarianceMatrixAdaptation(2)

    metrics = {'fitness_mean': [], 'fitness_max': [], 'c_norm': []}

    for gen in range(1, 20):
        ranked_results, info = n_cma.step(spike_one)
        selected = torch.stack([g['parameters'] for g in ranked_results[0:n_cma.mu]])
        d, b = torch.linalg.eigh(n_cma.c)
        d = d.sqrt().diag_embed()
        plot_heatmap('select', gen, n_cma.mean, b, d, g=selected)
        for key in metrics:
            metrics[key].append(info[key])

    for key in metrics:
        plt.title(key)
        plt.plot(metrics[key])
        plt.legend(loc='lower left')
        plt.show()
//  [python]
def renyi_entropy(self, x):        
        k = self.kernel(x)
        k = k / torch.trace(k)
        eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
        eig_pow = eigv ** alpha
        entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
        return entropy
//  [python]
def renyi_entropy(self, x):        
        k = self.kernel(x)
        k = k / torch.trace(k)
        eigv = torch.abs(torch.linalg.eigh(k)[0])
        eig_pow = eigv ** alpha
        entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
        return entropy
//  [python]
def joint_entropy(self, x, y):
        x = self.kernel(x)
        y = self.kernel(y)
        k = torch.mul(x, y)
        k = k / torch.trace(k)
        eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
        eig_pow =  eigv ** self.alpha
        entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
        return entropy
//  [python]
def joint_entropy(self, x, y):
        x = self.kernel(x)
        y = self.kernel(y)
        k = torch.mul(x, y)
        k = k / torch.trace(k)
        eigv = torch.abs(torch.linalg.eigh(k)[0])
        eig_pow =  eigv ** self.alpha
        entropy = (1 / (1 - self.alpha)) * torch.log2(torch.sum(eig_pow))
        return entropy
//  [python]
def reyi_entropy(x, sigma):
    alpha = 1.01
    k = calculate_gram_mat(x, sigma)
    k = k / torch.trace(k)
    eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
    eig_pow = eigv ** alpha
    entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def reyi_entropy(x, sigma):
    alpha = 1.01
    k = calculate_gram_mat(x, sigma)
    k = k / torch.trace(k)
    eigv = torch.abs(torch.linalg.eigh(k)[0])
    eig_pow = eigv ** alpha
    entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def test_symeig():
    torch.manual_seed(100)
    X = torch.randn(1000, 5000)
    M = X @ X.T / 5000
    Ltr, Wtr = torch.symeig(M, eigenvectors=True)
    L, W = symeig(M)
    assert torch.abs(L - Ltr).max() / Ltr.max() < 1e-5
    rectr = Wtr @ torch.diag(Ltr) @ Wtr.T
    rec = W @ torch.diag(L) @ W.T
    assert torch.allclose(rectr, rec)
    assert torch.allclose(rec, M, rtol=1e-5, atol=1e-5)
//  [python]
def test_symeig():
    torch.manual_seed(100)
    X = torch.randn(1000, 5000)
    M = X @ X.T / 5000
    Ltr, Wtr = torch.linalg.eigh(M)
    L, W = symeig(M)
    assert torch.abs(L - Ltr).max() / Ltr.max() < 1e-5
    rectr = Wtr @ torch.diag(Ltr) @ Wtr.T
    rec = W @ torch.diag(L) @ W.T
    assert torch.allclose(rectr, rec)
    assert torch.allclose(rec, M, rtol=1e-5, atol=1e-5)
//  [python]
def get_process_eigval(S, Sigma, state1, state2, epsilon=1e-6, mode='regularize'):  
    ''' state can be int or list of int'''
    Sigma_sqrt_inv, Sigma_sqrt = sym_inverse(Sigma, epsilon, return_sqrt=True, mode=mode, return_both=True)
    
    S_similar = Sigma_sqrt @ S @ Sigma_sqrt

    eigval_all, eigvec_all = torch.symeig(S_similar, eigenvectors=True)
    eigvecs_K = Sigma_sqrt_inv @ eigvec_all

    # Find the relevant process which is changing most between state1 and state2
    process_id = torch.argmin(eigvecs_K[state1,:]*eigvecs_K[state2,:], dim=1).detach()
    
    return eigval_all[process_id]
//  [python]
def get_process_eigval(S, Sigma, state1, state2, epsilon=1e-6, mode='regularize'):  
    ''' state can be int or list of int'''
    Sigma_sqrt_inv, Sigma_sqrt = sym_inverse(Sigma, epsilon, return_sqrt=True, mode=mode, return_both=True)
    
    S_similar = Sigma_sqrt @ S @ Sigma_sqrt

    eigval_all, eigvec_all = torch.linalg.eigh(S_similar)
    eigvecs_K = Sigma_sqrt_inv @ eigvec_all

    # Find the relevant process which is changing most between state1 and state2
    process_id = torch.argmin(eigvecs_K[state1,:]*eigvecs_K[state2,:], dim=1).detach()
    
    return eigval_all[process_id]
//  [python]
def forward(self, A, X):
        with torch.autograd.set_detect_anomaly(True):

            lr = A
            lr_dim = lr.shape[0]
            f = X
            eig_val_lr, U_lr = torch.symeig(lr, eigenvectors=True, upper=True)

            # U_lr = torch.abs(U_lr)
            eye_mat = torch.eye(lr_dim).type(torch.FloatTensor)
            s_d = torch.cat((eye_mat, eye_mat), 0)

            a = torch.matmul(self.weights, s_d)
            b = torch.matmul(a, torch.t(U_lr))
            f_d = torch.matmul(b, f)
            f_d = torch.abs(f_d)
            f_d = f_d.fill_diagonal_(1)
            adj = f_d

            X = torch.mm(adj, adj.t())
            X = (X + X.t())/2
            X = X.fill_diagonal_(1)
        return adj, torch.abs(X)
//  [python]
def forward(self, A, X):
        with torch.autograd.set_detect_anomaly(True):

            lr = A
            lr_dim = lr.shape[0]
            f = X
            eig_val_lr, U_lr = torch.linalg.eigh(lr)

            # U_lr = torch.abs(U_lr)
            eye_mat = torch.eye(lr_dim).type(torch.FloatTensor)
            s_d = torch.cat((eye_mat, eye_mat), 0)

            a = torch.matmul(self.weights, s_d)
            b = torch.matmul(a, torch.t(U_lr))
            f_d = torch.matmul(b, f)
            f_d = torch.abs(f_d)
            f_d = f_d.fill_diagonal_(1)
            adj = f_d

            X = torch.mm(adj, adj.t())
            X = (X + X.t())/2
            X = X.fill_diagonal_(1)
        return adj, torch.abs(X)
//  [python]
def joint_entropy(x, y, s_x, s_y):
    alpha = 1.01
    x = calculate_gram_mat(x, s_x)
    y = calculate_gram_mat(y, s_y)
    k = torch.mul(x, y)
    k = k / torch.trace(k)
    eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
    eig_pow = eigv ** alpha
    entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def joint_entropy(x, y, s_x, s_y):
    alpha = 1.01
    x = calculate_gram_mat(x, s_x)
    y = calculate_gram_mat(y, s_y)
    k = torch.mul(x, y)
    k = k / torch.trace(k)
    eigv = torch.abs(torch.linalg.eigh(k)[0])
    eig_pow = eigv ** alpha
    entropy = (1 / (1 - alpha)) * torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def ExpH(a):
    # v0.3 OK
    """
    This function performs

            :math:`e^{H}`

    where H is the hermitian matrix.
    The Intricate computation follows procedure: symeig() -> exp() the singular matrix.

    Args:

        a :
            UniTensor, Must be a rank-2 [untagged], with one inbond, one outbond. If pass a non-rank2 tensor, tagged tensor, or pass a non-hermitian rank2 tensor; it will raise Error.

    Return:

        UniTensor, rank-2 [unregular], same bonds and labels and braket form as the original H
    """

    if isinstance(a,UniTensor):
        
        if a.is_symm:
           raise Exception("ExpH(a)","don't support symm tensor. GetBlock first.")

        else:
            if a.braket is not None:
                raise Exception("ExpH(a)","can only accept [untagged] type UniTensor")

            if a.is_diag:   
                u = torch.exp(a.Storage)
                
                tmp =  UniTensor(bonds=a.bonds,\
                                 labels=a.labels,\
                                 rowrank=a.rowrank,\
                                 is_diag=True,\
                                 check=False)
                tmp._mac(torch_tensor = u)
                return tmp
            else:
                if a.rowrank != 1:
                    raise Exception("ExpH(a)","a should be rank-2 tensor with 1 inbond 1 outbond")

                ## version-1, only real, not sure if it can extend to complex
                s , u = torch.symeig(a.Storage,eigenvectors=True)
                s     = torch.exp(s)

                # torch.matmul(u*s,u.transpose(0,1),out=u)
                u = torch.matmul(u*s,u.transpose(0,1))
                del s

                tmp =  UniTensor(bonds=a.bonds,\
                                labels=a.labels,\
                                rowrank=a.rowrank,\
                                check=False)
                tmp._mac(torch_tensor = u)
                return tmp

    else:
        raise Exception("ExpH(UniTensor)","[ERROR] ExpH can only accept UniTensor")
//  [python]
def ExpH(a):
    # v0.3 OK
    """
    This function performs

            :math:`e^{H}`

    where H is the hermitian matrix.
    The Intricate computation follows procedure: symeig() -> exp() the singular matrix.

    Args:

        a :
            UniTensor, Must be a rank-2 [untagged], with one inbond, one outbond. If pass a non-rank2 tensor, tagged tensor, or pass a non-hermitian rank2 tensor; it will raise Error.

    Return:

        UniTensor, rank-2 [unregular], same bonds and labels and braket form as the original H
    """

    if isinstance(a,UniTensor):
        
        if a.is_symm:
           raise Exception("ExpH(a)","don't support symm tensor. GetBlock first.")

        else:
            if a.braket is not None:
                raise Exception("ExpH(a)","can only accept [untagged] type UniTensor")

            if a.is_diag:   
                u = torch.exp(a.Storage)
                
                tmp =  UniTensor(bonds=a.bonds,\
                                 labels=a.labels,\
                                 rowrank=a.rowrank,\
                                 is_diag=True,\
                                 check=False)
                tmp._mac(torch_tensor = u)
                return tmp
            else:
                if a.rowrank != 1:
                    raise Exception("ExpH(a)","a should be rank-2 tensor with 1 inbond 1 outbond")

                ## version-1, only real, not sure if it can extend to complex
                s , u = torch.linalg.eigh(a.Storage)
                s     = torch.exp(s)

                # torch.matmul(u*s,u.transpose(0,1),out=u)
                u = torch.matmul(u*s,u.transpose(0,1))
                del s

                tmp =  UniTensor(bonds=a.bonds,\
                                labels=a.labels,\
                                rowrank=a.rowrank,\
                                check=False)
                tmp._mac(torch_tensor = u)
                return tmp

    else:
        raise Exception("ExpH(UniTensor)","[ERROR] ExpH can only accept UniTensor")
//  [python]
def _full_eigen(self, model, x):
        batch_size = x.shape[0]
        assert batch_size==1
        x = Variable(x, requires_grad=True)
        logits = forward(model, x)
        self.max_prob, self.max_index = F.softmax(logits, dim=1).max(1)
        y = logits.max(1)[1]
        loss = F.cross_entropy(logits, y)
        x_grad, = torch.autograd.grad(loss, x)
        x_grad = x_grad.view(batch_size, -1)

        W = self._get_W(x, model)
        probs = F.softmax(logits, 1)
        W = W.cpu().detach()
        probs = probs.cpu().detach()

        D = torch.diag(probs[0, :])
        A = (D - probs.t().mm(probs))
        sigma_A, U_A = torch.symeig(-A, eigenvectors=True)
        rank_A = torch.sum(sigma_A<0)
        sigma_A = -sigma_A[:rank_A]
        U_A = U_A[:, :rank_A]

        sigma_A_sqrt = torch.sqrt(sigma_A)
        sigma_A_sqrt = torch.diag(sigma_A_sqrt)
        chol_A = U_A.mm(sigma_A_sqrt)
        B = W.mm(chol_A)
        BTB = B.t().mm(B)
        rank_B = np.linalg.matrix_rank(BTB)

        sigma_B_sq, V_B = torch.symeig(-BTB, eigenvectors=True)
        sigma_B_sq = -sigma_B_sq[:rank_B]
        V_B = V_B[:, :rank_B]

        sigma_B_inv = torch.rsqrt(sigma_B_sq)        
        sigma_B_inv = torch.diag(sigma_B_inv)
        HEV = V_B.mm(sigma_B_inv)
        HEV = B.mm(HEV)

        HEV = HEV.type(torch.cuda.FloatTensor)
        W = W.type(torch.cuda.FloatTensor)
        A = A.type(torch.cuda.FloatTensor)
        hvp_fn = lambda z: ((z.mm(W)).mm(A)).mm(W.t())
        x_grad_data = x_grad.detach()
        sigma_H = sigma_B_sq.detach()
            
        return x_grad_data, sigma_H, HEV, hvp_fn
//  [python]
def _full_eigen(self, model, x):
        batch_size = x.shape[0]
        assert batch_size==1
        x = Variable(x, requires_grad=True)
        logits = forward(model, x)
        self.max_prob, self.max_index = F.softmax(logits, dim=1).max(1)
        y = logits.max(1)[1]
        loss = F.cross_entropy(logits, y)
        x_grad, = torch.autograd.grad(loss, x)
        x_grad = x_grad.view(batch_size, -1)

        W = self._get_W(x, model)
        probs = F.softmax(logits, 1)
        W = W.cpu().detach()
        probs = probs.cpu().detach()

        D = torch.diag(probs[0, :])
        A = (D - probs.t().mm(probs))
        sigma_A, U_A = torch.linalg.eigh(-A)
        rank_A = torch.sum(sigma_A<0)
        sigma_A = -sigma_A[:rank_A]
        U_A = U_A[:, :rank_A]

        sigma_A_sqrt = torch.sqrt(sigma_A)
        sigma_A_sqrt = torch.diag(sigma_A_sqrt)
        chol_A = U_A.mm(sigma_A_sqrt)
        B = W.mm(chol_A)
        BTB = B.t().mm(B)
        rank_B = np.linalg.matrix_rank(BTB)

        sigma_B_sq, V_B = torch.linalg.eigh(-BTB)
        sigma_B_sq = -sigma_B_sq[:rank_B]
        V_B = V_B[:, :rank_B]

        sigma_B_inv = torch.rsqrt(sigma_B_sq)        
        sigma_B_inv = torch.diag(sigma_B_inv)
        HEV = V_B.mm(sigma_B_inv)
        HEV = B.mm(HEV)

        HEV = HEV.type(torch.cuda.FloatTensor)
        W = W.type(torch.cuda.FloatTensor)
        A = A.type(torch.cuda.FloatTensor)
        hvp_fn = lambda z: ((z.mm(W)).mm(A)).mm(W.t())
        x_grad_data = x_grad.detach()
        sigma_H = sigma_B_sq.detach()
            
        return x_grad_data, sigma_H, HEV, hvp_fn
//  [python]
def test_whitening_correctness3():
    '''
    This tests assess if whitening matrix can
    be correctly generated for a real covariance matrix
    '''
    from scipy.linalg import inv, sqrtm

    w = WhitenTraj("data/whitened")

    # Load three different covariance matrices
    # cov1 is a DiffNets-generated matrix for 8 myosin isoforms with a subset of atoms relative to cov2
    # cov2 is a DiffNets-generated matrix for 8 myosin isoforms
    # cov3 is beta-lactamase covariance matrix
    for i in range(1, 4):
        # Load covariance matrix
        cov = np.load(f"data/cov/cov{i}.npy")

        e, v = torch.symeig(torch.from_numpy(cov).double(), eigenvectors=True)
        # In valid covariance matrix the smallest eigenvalue should be positive
        # because the covariance matrix is a positive semidefinite matrix
        # https://stats.stackexchange.com/questions/52976/is-a-sample-covariance-matrix-always-symmetric-and-positive-definite
        assert torch.min(e) > 0.0

        # Covariance matrix should be symmetric
        assert np.allclose(cov, cov.T)

        uwm, wm = w.get_wuw_mats(cov.astype(np.double))

        # ZCA whitening matrix should be symmteric
        assert np.allclose(wm, wm.T)

        # assert that whitening and whitening produce the identity matrix
        assert_matrix_is_identity(np.matmul(wm, uwm))

        # assert that the covariance matrix multiplied by the transpose of
        # the whitening matrix multiplied by the whitening matrix is identity
        # for detailed explanation see page 3 of https://arxiv.org/pdf/1512.00809.pdf
        test = np.matmul(cov, np.matmul(wm.transpose(), wm))
        assert_matrix_is_identity(test)
//  [python]
def test_whitening_correctness3():
    '''
    This tests assess if whitening matrix can
    be correctly generated for a real covariance matrix
    '''
    from scipy.linalg import inv, sqrtm

    w = WhitenTraj("data/whitened")

    # Load three different covariance matrices
    # cov1 is a DiffNets-generated matrix for 8 myosin isoforms with a subset of atoms relative to cov2
    # cov2 is a DiffNets-generated matrix for 8 myosin isoforms
    # cov3 is beta-lactamase covariance matrix
    for i in range(1, 4):
        # Load covariance matrix
        cov = np.load(f"data/cov/cov{i}.npy")

        e, v = torch.linalg.eigh(torch.from_numpy(cov).double())
        # In valid covariance matrix the smallest eigenvalue should be positive
        # because the covariance matrix is a positive semidefinite matrix
        # https://stats.stackexchange.com/questions/52976/is-a-sample-covariance-matrix-always-symmetric-and-positive-definite
        assert torch.min(e) > 0.0

        # Covariance matrix should be symmetric
        assert np.allclose(cov, cov.T)

        uwm, wm = w.get_wuw_mats(cov.astype(np.double))

        # ZCA whitening matrix should be symmteric
        assert np.allclose(wm, wm.T)

        # assert that whitening and whitening produce the identity matrix
        assert_matrix_is_identity(np.matmul(wm, uwm))

        # assert that the covariance matrix multiplied by the transpose of
        # the whitening matrix multiplied by the whitening matrix is identity
        # for detailed explanation see page 3 of https://arxiv.org/pdf/1512.00809.pdf
        test = np.matmul(cov, np.matmul(wm.transpose(), wm))
        assert_matrix_is_identity(test)
//  [python]
def nasi(model, x, target, criterion):
    ntk = empirical_ntk(model, x)
    eigenvalues, _ = torch.symeig(ntk)  # ascending
    return np.sqrt(eigenvalues.cpu().numpy().sum() / x.shape[0])
//  [python]
def nasi(model, x, target, criterion):
    ntk = empirical_ntk(model, x)
    eigenvalues, _ = torch.linalg.eigh(ntk)
    return np.sqrt(eigenvalues.cpu().numpy().sum() / x.shape[0])
//  [python]
def eig(X):
    d = X.shape[-1]
    if d == 2:
        return eig2x2(X)
    # TODO(ccruceru): https://github.com/pytorch/pytorch/pull/22909.
    # elif d == 3:
    #     return eig3x3(X)

    # eigenvectors for backward
    return torch.symeig(X, eigenvectors=True).eigenvalues
//  [python]
def eig(X):
    d = X.shape[-1]
    if d == 2:
        return eig2x2(X)
    # TODO(ccruceru): https://github.com/pytorch/pytorch/pull/22909.
    # elif d == 3:
    #     return eig3x3(X)

    # eigenvectors for backward
    return torch.linalg.eigh(X).eigenvalues
//  [python]
def forward(ctx, a):
        vals, vecs = torch.symeig(a)
        ctx.save_for_backward(vals, vecs)
        return vecs @ vals.abs().sqrt().diag_embed() @ vecs.transpose(-2, -1)
//  [python]
def forward(ctx, a):
        vals, vecs = torch.linalg.eigh(a)
        ctx.save_for_backward(vals, vecs)
        return vecs @ vals.abs().sqrt().diag_embed() @ vecs.transpose(-2, -1)
//  [python]
def sqrtm_torch(x):
    """
    This function computes the square root of a matrix.

    Parameters
    ----------
    :param x: positive definite matrix

    Returns
    -------
    :return: sqrtm(x)
    """
    eigenvalues, eigenvectors = torch.symeig(x, eigenvectors=True)

    sqrt_eigenvalues = torch.sqrt(torch.maximum(eigenvalues, torch.zeros_like(eigenvalues)))  # Assume real eigenvalues (first column only)

    return torch.matmul(eigenvectors, torch.matmul(torch.diag_embed(sqrt_eigenvalues), torch.inverse(eigenvectors)))
//  [python]
def sqrtm_torch(x):
    """
    This function computes the square root of a matrix.

    Parameters
    ----------
    :param x: positive definite matrix

    Returns
    -------
    :return: sqrtm(x)
    """
    eigenvalues, eigenvectors = torch.linalg.eigh(x)

    sqrt_eigenvalues = torch.sqrt(torch.maximum(eigenvalues, torch.zeros_like(eigenvalues)))  # Assume real eigenvalues (first column only)

    return torch.matmul(eigenvectors, torch.matmul(torch.diag_embed(sqrt_eigenvalues), torch.inverse(eigenvectors)))
//  [python]
def eigh(tensor: TensorType) -> torch.Tensor:
        "Computes the eigen decomposition of a self-adjoint matrix."
        eigval, eigvec = torch.symeig(tensor)
        return eigval, eigvec
//  [python]
def eigh(tensor: TensorType) -> torch.Tensor:
        "Computes the eigen decomposition of a self-adjoint matrix."
        eigval, eigvec = torch.linalg.eigh(tensor)
        return eigval, eigvec
//  [python]
def matrix_sqrt(M, thresh=False):
    S, V = torch.eig(M.cuda())
    S = S[:,0]
    S[S<0] = 0

    if thresh:
        k = int(3*len(S)//4)
        print("Thresh k:",k)
        S[:k] = 0

    S = torch.diag(S**0.5)
    return (V @ S @ V.T).cpu()
//  [python]
def matrix_sqrt(M, thresh=False):
    S, V = torch.linalg.eigh(M.cuda())
    S[S<0] = 0

    if thresh:
        k = int(3*len(S)//4)
        print("Thresh k:",k)
        S[:k] = 0

    S = torch.diag(S**0.5)
    return (V @ S @ V.T).cpu()
//  [python]
def eig_coord_from_dist(D):
    M = (D[:1, :] + D[:, :1] - D) / 2
    L, V = torch.eig(M, eigenvectors=True)
    L = torch.diag_embed(torch.sort(L[:,0], descending=True)[0])
    X = torch.matmul(V, L.clamp(min=0).sqrt())
    return X[:, :3].detach()
//  [python]
def eig_coord_from_dist(D):
    M = (D[:1, :] + D[:, :1] - D) / 2
    L, V = torch.linalg.eigh(M)
    L = torch.diag_embed(torch.sort(L, descending=True)[0])
    X = torch.matmul(V, L.clamp(min=0).sqrt())
    return X[:, :3].detach()
//  [python]
def K(C_t):
    return (torch.symeig(C_t, eigenvectors=False).eigenvalues.min())
//  [python]
def K(C_t):
    return (torch.linalg.eigh(C_t).eigenvalues.min())
//  [python]
def P(C_t):
    return (torch.prod(torch.symeig(C_t).eigenvalues**(1/784)))
//  [python]
def P(C_t):
    return (torch.prod(linalg.eigh(C_t).eigenvalues**(1/784)))
//  [python]
def test_DominantSymeig():
    N = 300
    K = torch.randn(N, N, dtype=torch.float64)
    K = K + K.T
    target = torch.randn(N, dtype=torch.float64)
    potential = torch.randn(N, dtype=torch.float64)
    potential.requires_grad_(True)
    V = torch.diag(potential)
    H = K + V
    print("\n----- test_DominantSymeig -----")
    print("----- Dimension of the Hamiltonian: %d -----" % N)

    for i in range(3):
        ##### Backward through torch.symeig implemented by Pytorch.
        start1 = time.time()
        Es, psis = torch.symeig(H, eigenvectors=True)
        psi0_torch = psis[:, 0]
        loss_torch = 1. - torch.matmul(psi0_torch, target)
        end1 = time.time()

        start2 = time.time()
        grad_potential_torch, = torch.autograd.grad(loss_torch, potential)
        end2 = time.time()
        print("Pytorch\t\tforward time: %f    backward time: %f" % (end1 - start1, end2 - start2))

        ##### Backward through the DominantSymeig primitive.
        dominant_symeig = DominantSymeig.apply
        k = 300
        start1 = time.time()
        _, psi0_dominant = dominant_symeig(H, k)
        loss_dominant = 1. - torch.matmul(psi0_dominant, target)
        end1 = time.time()

        start2 = time.time()
        grad_potential_dominant, = torch.autograd.grad(loss_dominant, potential)
        end2 = time.time()
        print("DominantSymeig\tforward time: %f    backward time: %f" % (end1 - start1, end2 - start2))

        assert torch.allclose(loss_dominant, loss_torch) or \
               torch.allclose(loss_dominant, 2. - loss_torch)
        assert torch.allclose(grad_potential_dominant, grad_potential_torch) or \
               torch.allclose(grad_potential_dominant, -grad_potential_torch)
//  [python]
def test_DominantSymeig():
    N = 300
    K = torch.randn(N, N, dtype=torch.float64)
    K = K + K.T
    target = torch.randn(N, dtype=torch.float64)
    potential = torch.randn(N, dtype=torch.float64)
    potential.requires_grad_(True)
    V = torch.diag(potential)
    H = K + V
    print("\n----- test_DominantSymeig -----")
    print("----- Dimension of the Hamiltonian: %d -----" % N)

    for i in range(3):
        ##### Backward through torch.linalg.eigh implemented by Pytorch.
        start1 = time.time()
        Es, psis = torch.linalg.eigh(H)
        psi0_torch = psis[:, 0]
        loss_torch = 1. - torch.matmul(psi0_torch, target)
        end1 = time.time()

        start2 = time.time()
        grad_potential_torch, = torch.autograd.grad(loss_torch, potential)
        end2 = time.time()
        print("Pytorch\t\tforward time: %f    backward time: %f" % (end1 - start1, end2 - start2))

        ##### Backward through the DominantSymeig primitive.
        dominant_symeig = DominantSymeig.apply
        k = 300
        start1 = time.time()
        _, psi0_dominant = dominant_symeig(H, k)
        loss_dominant = 1. - torch.matmul(psi0_dominant, target)
        end1 = time.time()

        start2 = time.time()
        grad_potential_dominant, = torch.autograd.grad(loss_dominant, potential)
        end2 = time.time()
        print("DominantSymeig\tforward time: %f    backward time: %f" % (end1 - start1, end2 - start2))

        assert torch.allclose(loss_dominant, loss_torch) or \
               torch.allclose(loss_dominant, 2. - loss_torch)
        assert torch.allclose(grad_potential_dominant, grad_potential_torch) or \
               torch.allclose(grad_potential_dominant, -grad_potential_torch)
//  [python]
def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [[rot[..., i, j] for j in range(3)] for i in range(3)]

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    k_tensor = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    evals, vectors = torch.linalg.eig(k_tensor)
    
    max_idx = torch.argmax(evals.real, dim=-1)
    
    vectors = vectors.gather(dim=-1,index=max_idx[...,None].expand(-1,-1,4))
    
    return vectors[...,0]
//  [python]
def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [[rot[..., i, j] for j in range(3)] for i in range(3)]

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    _, vectors = torch.linalg.eigh((1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2))
    return vectors[..., -1]
//  [python]
def __init__(self, matrix) -> None:
        eigen_value, eigen_vector = torch.symeig(matrix, eigenvectors=True)
        self.value = eigen_value
        self.vector = eigen_vector
//  [python]
def __init__(self, matrix) -> None:
        eigen_value, eigen_vector = torch.linalg.eigh(matrix, UPLO='U')
        self.value = eigen_value
        self.vector = eigen_vector
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if(rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot 

    k = [
        [ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
        [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
        [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
        [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]
    ]

    k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if(rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot 

    k = [
        [ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
        [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
        [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
        [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]
    ]

    k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def compute_rotation_QCQP_grad_fast(A, E, nu, x):
    """
    Input: A: (B,10,10) tensor (parametrices B symmetric 4x4 matrices)
           E: (22,10,10) tensor (quadratic symmetric equality constraint matrices)
           nu: (B,22) tensor (optimal lagrange multipliers)
           x: (B,10) tensor (optimal vectorized rotation matrices with homogenizing 10th entry of 1)

    Output: grad: (B, 10, 55) tensor (gradient)

    Applies the implicit function theorem to compute gradients of the solution to an equality-constrained
    homogeneous rotation matrix QCQP.
    """
    assert(A.dim() > 2)
    assert(E.dim() > 2)
    assert(nu.dim() > 1)
    assert(x.dim() > 1)

    # Remove redundant/SO(3) constraints
    num_constraints = 7
    M = A.new_zeros((A.shape[0], 10 + num_constraints, 10 + num_constraints))

    # TODO: are all 22 E's needed, or just the 7 of interest? Should be all 22 for KKT definitino
    M[:, :10, :10] = A + torch.einsum('bi,imn->bmn', nu, E)
    B = torch.einsum('mij,bj->bim', E[:num_constraints, :, :], x)
    M[:, :10, 10:] = B
    M[:, 10:, :10] = B.transpose(1, 2)

    # Eig check
    eigs, _ = torch.symeig(M)

    b = A.new_zeros((A.shape[0], 10+num_constraints, 55))
    # symmetric matrix indices
    idx = torch.triu_indices(10, 10)

    i = torch.arange(55)
    I_ij = A.new_zeros((55, 10, 10))

    I_ij[i, idx[0], idx[1]] = 1.
    I_ij[i, idx[1], idx[0]] = 1.

    I_ij = I_ij.expand(A.shape[0], 55, 10, 10)

    b[:, :10, :] = torch.einsum('bkij,bi->bjk', I_ij, x)

    # This solves all gradients simultaneously!
    X, _ = torch.solve(b, M)

    grad = -1 * X[:, :10, :]

    return grad
//  [python]
def compute_rotation_QCQP_grad_fast(A, E, nu, x):
    """
    Input: A: (B,10,10) tensor (parametrices B symmetric 4x4 matrices)
           E: (22,10,10) tensor (quadratic symmetric equality constraint matrices)
           nu: (B,22) tensor (optimal lagrange multipliers)
           x: (B,10) tensor (optimal vectorized rotation matrices with homogenizing 10th entry of 1)

    Output: grad: (B, 10, 55) tensor (gradient)

    Applies the implicit function theorem to compute gradients of the solution to an equality-constrained
    homogeneous rotation matrix QCQP.
    """
    assert(A.dim() > 2)
    assert(E.dim() > 2)
    assert(nu.dim() > 1)
    assert(x.dim() > 1)

    # Remove redundant/SO(3) constraints
    num_constraints = 7
    M = A.new_zeros((A.shape[0], 10 + num_constraints, 10 + num_constraints))

    # TODO: are all 22 E's needed, or just the 7 of interest? Should be all 22 for KKT definitino
    M[:, :10, :10] = A + torch.einsum('bi,imn->bmn', nu, E)
    B = torch.einsum('mij,bj->bim', E[:num_constraints, :, :], x)
    M[:, :10, 10:] = B
    M[:, 10:, :10] = B.transpose(1, 2)

    # Eig check
    eigs, _ = torch.linalg.eigh(M)

    b = A.new_zeros((A.shape[0], 10+num_constraints, 55))
    # symmetric matrix indices
    idx = torch.triu_indices(10, 10)

    i = torch.arange(55)
    I_ij = A.new_zeros((55, 10, 10))

    I_ij[i, idx[0], idx[1]] = 1.
    I_ij[i, idx[1], idx[0]] = 1.

    I_ij = I_ij.expand(A.shape[0], 55, 10, 10)

    b[:, :10, :] = torch.einsum('bkij,bi->bjk', I_ij, x)

    # This solves all gradients simultaneously!
    X, _ = torch.solve(b, M)

    grad = -1 * X[:, :10, :]

    return grad
//  [python]
def rotmat2qvec(R: torch.Tensor):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flatten()
    K = torch.Tensor([[Rxx - Ryy - Rzz, 0, 0, 0],
                      [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
                      [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
                      [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
    eigvals, eigvecs = torch.eig(K, eigenvectors=True)
    qvec = eigvecs[:, torch.argmax(eigvals.real)][[3, 0, 1, 2]]
    if qvec[0] < 0:
        qvec *= -1
    return qvec
//  [python]
def rotmat2qvec(R: torch.Tensor):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flatten()
    K = torch.Tensor([[Rxx - Ryy - Rzz, 0, 0, 0],
                      [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
                      [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
                      [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
    eigvals, eigvecs = torch.linalg.eigh(K)
    qvec = eigvecs[[3, 0, 1, 2], torch.argmax(eigvals)]
    if qvec[0] < 0:
        qvec *= -1
    return qvec
//  [python]
def forward(ctx, A):
        eival, eivec = torch.symeig(A)
        ctx.save_for_backward(eival, eivec)
        return eival, eivec
//  [python]
def forward(ctx, A):
        eival, eivec = torch.linalg.eigh(A)
        ctx.save_for_backward(eival, eivec)
        return eival, eivec
//  [python]
def get_whitening_parameters(patches):
    n,c,h,w = patches.shape
    est_covariance = torch.cov(patches.view(n, c*h*w).t())
    eigenvalues, eigenvectors = torch.symeig(est_covariance, eigenvectors=True)
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.t().reshape(c*h*w,c,h,w).flip(0)
//  [python]
def get_whitening_parameters(patches):
    # As a high-level summary, we're basically finding the high-dimensional oval that best fits the data here.
    # We can then later use this information to map the input information to a nicely distributed sphere, where also
    # the most significant features of the inputs each have their own axis. This significantly cleans things up for the
    # rest of the neural network and speeds up training.
    n,c,h,w = patches.shape
    est_covariance = torch.cov(patches.view(n, c*h*w).t())
    eigenvalues, eigenvectors = torch.linalg.eigh(est_covariance, UPLO='U') # this is the same as saying we want our eigenvectors, with the specification that the matrix be an upper triangular matrix (instead of a lower-triangular matrix)
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.t().reshape(c*h*w,c,h,w).flip(0)
//  [python]
def rot_to_quat(
        rot: torch.Tensor,
):
    if (rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [xx + yy + zz, zy - yz, xz - zx, yx - xy, ],
        [zy - yz, xx - yy - zz, xy + yx, xz + zx, ],
        [xz - zx, xy + yx, yy - xx - zz, yz + zy, ],
        [yx - xy, xz + zx, yz + zy, zz - xx - yy, ]
    ]

    k = (1. / 3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    evals, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(
        rot: torch.Tensor,
):
    if (rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [xx + yy + zz, zy - yz, xz - zx, yx - xy, ],
        [zy - yz, xx - yy - zz, xy + yx, xz + zx, ],
        [xz - zx, xy + yx, yy - xx - zz, yz + zy, ],
        [yx - xy, xz + zx, yz + zy, zz - xx - yy, ]
    ]

    k = (1. / 3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def _get_eigenvectors(Eyyt: torch.Tensor, num_data_steps: int) -> torch.Tensor:
    Eyyt = Eyyt / num_data_steps
    # https://stats.stackexchange.com/questions/390532/adding-a-small-constant-to-the-diagonals-of-a-matrix-to-stabilize
    damp = EIGEN_DAMPEN_FACTOR * torch.mean(torch.diag(Eyyt))
    diag = torch.arange(Eyyt.shape[-1], device=Eyyt.device)
    Eyyt[diag, diag] = Eyyt[diag, diag] + damp
    logger.info(f"Covariance matrix dtype before eigh {Eyyt.dtype=}")
    _, u = torch.symeig(Eyyt)
    return u
//  [python]
def _get_eigenvectors(Eyyt: torch.Tensor, num_data_steps: int) -> torch.Tensor:
    Eyyt = Eyyt / num_data_steps
    # https://stats.stackexchange.com/questions/390532/adding-a-small-constant-to-the-diagonals-of-a-matrix-to-stabilize
    damp = EIGEN_DAMPEN_FACTOR * torch.mean(torch.diag(Eyyt))
    diag = torch.arange(Eyyt.shape[-1], device=Eyyt.device)
    Eyyt[diag, diag] = Eyyt[diag, diag] + damp
    logger.info(f"Covariance matrix dtype before eigh {Eyyt.dtype=}")
    _, u = torch.linalg.eigh(Eyyt)
    return u
//  [python]
def __init__(self, N = 256, a =3/(3.8**2)):
        J = torch.zeros(N, N)
        for i, j in zip(np.arange(N-1), np.arange(1, N)):
            J[i,i] += a
            J[j,j] += a
            J[i,j] = J[j,i] = -a
        D, P = torch.symeig(J)
        D_inv = 1/D
        D_inv[0] = 0
        self.P, self.D_inv = P, D_inv
        self.N = N
//  [python]
def __init__(self, N = 256, a =3/(3.8**2)):
        J = torch.zeros(N, N)
        for i, j in zip(np.arange(N-1), np.arange(1, N)):
            J[i,i] += a
            J[j,j] += a
            J[i,j] = J[j,i] = -a
        D, P = torch.linalg.eigh(J)
        D_inv = 1/D
        D_inv[0] = 0
        self.P, self.D_inv = P, D_inv
        self.N = N
//  [python]
def get_svqb(U, edge):
    tau = 1.2e-07
    UBU = U.T @ B(U, edge)
    d = UBU.diagonal(0, -2, -1)
    nz = torch.where(abs(d) != 0.0)
    assert len(nz) == 1, nz
    if len(nz[0]) < len(d):
        U = U[:, nz[0]]
        UBU = U.T @ B(U, edge)
        d = UBU.diagonal(0, -2, -1)
        nz = torch.where(abs(d) != 0.0)
        assert len(nz[0]) == len(d)
    d_col = (d ** -0.5).reshape(d.shape[0], 1)
    DUBUD = (UBU * d_col) * d_col.T
    E, Z = torch.symeig(DUBUD, eigenvectors=True, upper=True)
    t = tau * abs(E).max()
    keep = torch.where(E > t)
    assert len(keep) == 1, keep
    E = E[keep[0]]
    Z = Z[:, keep[0]]
    # check_shape(U, d_col, Z, E, (U * d_col.T), (Z * E ** -0.5))
    U = (U * d_col.T) @ (Z * E ** -0.5)
    UAU = U.T @ A(U, edge)
    E, Z = torch.symeig(UAU, eigenvectors=True, upper=True)
    return U @ Z[:, :n ], E[:n]
//  [python]
def get_svqb(U, edge):
    tau = 1.2e-07
    UBU = U.T @ B(U, edge)
    d = UBU.diagonal(0, -2, -1)
    nz = torch.where(abs(d) != 0.0)
    assert len(nz) == 1, nz
    if len(nz[0]) < len(d):
        U = U[:, nz[0]]
        UBU = U.T @ B(U, edge)
        d = UBU.diagonal(0, -2, -1)
        nz = torch.where(abs(d) != 0.0)
        assert len(nz[0]) == len(d)
    d_col = (d ** -0.5).reshape(d.shape[0], 1)
    DUBUD = (UBU * d_col) * d_col.T
    E, Z = torch.linalg.eigh(DUBUD, UPLO='U')
    t = tau * abs(E).max()
    keep = torch.where(E > t)
    assert len(keep) == 1, keep
    E = E[keep[0]]
    Z = Z[:, keep[0]]
    # check_shape(U, d_col, Z, E, (U * d_col.T), (Z * E ** -0.5))
    U = (U * d_col.T) @ (Z * E ** -0.5)
    UAU = U.T @ A(U, edge)
    E, Z = torch.linalg.eigh(UAU, UPLO='U')
    return U @ Z[:, :n ], E[:n]
//  [python]
def positivize(a: Tensor, eps: float = 1e-3) -> Tensor:
    assert a.ndim == 2
    w, v = torch.eig(a, eigenvectors=True)
    w = w[:, 0].clamp(min=eps).sqrt()
    sqrt = v.mm(torch.diag(w))
    a = sqrt.mm(sqrt.t())
    return a
//  [python]
def positivize(a: Tensor, eps: float = 1e-3) -> Tensor:
    assert a.ndim == 2
    w, v = torch.linalg.eigh(a)
    s = w.clamp(min=eps).sqrt()
    sqrt = v.mm(s.diag_embed())
    a = sqrt.mm(sqrt.t())
    return a
//  [python]
def symeig(M):
    """Symetric eigendecomposition avoiding failure cases by
    adding and removing jitter to the diagonal.

    Parameters
    ----------
    M : torch.Tensor

    Returns
    -------
    L : torch.Tensor
        eigenvalues
    W : torch.Tensor
        eigenvectors
    """
    try:
        L, W = torch.symeig(M, eigenvectors=True, upper=True)
    except RuntimeError:  # did not converge
        logging.info('SYMEIG: adding jitter, did not converge.')
        # use W L W^T + I = W (L + I) W^T
        M = M + torch.eye(M.shape[0], device=M.device)
        try:
            L, W = torch.symeig(M, eigenvectors=True, upper=True)
            L -= 1.
        except RuntimeError:
            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
            exit()
    # eigenvalues of symeig at least 0
    L = L.clamp(min=0.0)
    L = torch.nan_to_num(L)
    W = torch.nan_to_num(W)
    return L, W
//  [python]
def symeig(M):
    """Symetric eigendecomposition avoiding failure cases by
    adding and removing jitter to the diagonal.

    Parameters
    ----------
    M : torch.Tensor

    Returns
    -------
    L : torch.Tensor
        eigenvalues
    W : torch.Tensor
        eigenvectors
    """
    try:
        L, W = torch.linalg.eigh(M, UPLO='U')
    except RuntimeError:  # did not converge
        logging.info('SYMEIG: adding jitter, did not converge.')
        # use W L W^T + I = W (L + I) W^T
        M = M + torch.eye(M.shape[0], device=M.device)
        try:
            L, W = torch.linalg.eigh(M, UPLO='U')
            L -= 1.
        except RuntimeError:
            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
            exit()
    # eigenvalues of symeig at least 0
    L = L.clamp(min=0.0)
    L = torch.nan_to_num(L)
    W = torch.nan_to_num(W)
    return L, W
//  [python]
def fit(K:Tensor, y:Tensor, train_mask:Tensor, nugget:Union[float, Tensor]=1e-2):
    """
    Make predictions for a range of nugget values in one shot.
    For classification problems, the prediction target is classification probability.
    For regression problems, the prediction target is output value.

    Args:
        K (Tensor[float]): the GNNGP Kernel matrix.
        y (Tensor[float] or Tensor[int]): prediction target.
            For int type, the task is a classification problem.
            For float type, the task is a regression problem.
        train_mask (Tensor[bool]): mask for training nodes.
        nugget (float or Tensor[float]): the nugget used in posterior inference.
            (default: 1e-2 x mean diagonal element of K)
    """
    mb = train_mask
    if isinstance(nugget, float): nugget = torch.tensor([nugget])
    if y.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
        yb = torch.nn.functional.one_hot(y[mb]).to(torch.float)
    else: yb = y[mb]
    fitted = torch.zeros((len(nugget), K.shape[0], *yb.shape[1:]), device=K.device)
    w, v = torch.symeig(K[mb][:,mb], eigenvectors=True)
    d = torch.mean(torch.diag(K))
    shape = (-1, 1) if len(yb.shape) >= 2 else (-1,)
    temp = v.T @ yb
    for i, eps in enumerate(nugget):
        fitted[i] = K[:,mb] @ (v @ (temp/(w+eps*d).view(shape)))
    return fitted[0] if len(nugget) == 1 else fitted
//  [python]
def fit(K:Tensor, y:Tensor, train_mask:Tensor, nugget:Union[float, Tensor]=1e-2):
    """
    Make predictions for a range of nugget values in one shot.
    For classification problems, the prediction target is classification probability.
    For regression problems, the prediction target is output value.

    Args:
        K (Tensor[float]): the GNNGP Kernel matrix.
        y (Tensor[float] or Tensor[int]): prediction target.
            For int type, the task is a classification problem.
            For float type, the task is a regression problem.
        train_mask (Tensor[bool]): mask for training nodes.
        nugget (float or Tensor[float]): the nugget used in posterior inference.
            (default: 1e-2 x mean diagonal element of K)
    """
    mb = train_mask
    if isinstance(nugget, float): nugget = torch.tensor([nugget])
    if y.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
        yb = torch.nn.functional.one_hot(y[mb]).to(torch.float)
    else: yb = y[mb]
    fitted = torch.zeros((len(nugget), K.shape[0], *yb.shape[1:]), device=K.device)
    w, v = torch.linalg.eigh(K[mb][:,mb])
    d = torch.mean(torch.diag(K))
    shape = (-1, 1) if len(yb.shape) >= 2 else (-1,)
    temp = v.T @ yb
    for i, eps in enumerate(nugget):
        fitted[i] = K[:,mb] @ (v @ (temp/(w+eps*d).view(shape)))
    return fitted[0] if len(nugget) == 1 else fitted