jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def initialization(self, n_items, embedding_dim, device=None):
        X = torch.randn((n_items, embedding_dim), device=device)
        X -= X.mean(axis=0)
        if not isinstance(n_items, torch.Tensor):
            n_items = torch.tensor(n_items, dtype=X.dtype, device=device)
        lmbda, Q = torch.symeig(X.T @ X, eigenvectors=True)
        X = X @ Q @ torch.diag(lmbda ** (-0.5)) * (n_items.type(X.dtype).sqrt())
        return X
//  [python]
def initialization(self, n_items, embedding_dim, device=None):
        X = torch.randn((n_items, embedding_dim), device=device)
        X -= X.mean(axis=0)
        if not isinstance(n_items, torch.Tensor):
            n_items = torch.tensor(n_items, dtype=X.dtype, device=device)
        lmbda, Q = torch.linalg.eigh(X.T @ X, UPLO="U")
        X = X @ Q @ torch.diag(lmbda ** (-0.5)) * (n_items.type(X.dtype).sqrt())
        return X
//  [python]
def _get_psd_logm(mat, method):
    if isinstance(mat, torch.Tensor):
        eps = torch.tensor(torch.finfo(mat.dtype).eps)
        if method=='eigen':
            EVL,EVC = torch.symeig(mat)
            ret = (EVC * torch.log(torch.maximum(eps, EVL))) @ EVC.T.conj()
        else:
            assert (len(method)==3) and (method[0]=='pade')
            logm_op = numqi._torch_op.get_PSDMatrixLogm(int(method[1]), int(method[2]))
            ret = logm_op(mat)
    else: #numpy
        eps = np.finfo(mat.dtype).eps
        EVL,EVC = np.linalg.eigh(mat)
        ret = (EVC * np.log(np.maximum(eps, EVL))) @ EVC.T.conj()
    return ret
//  [python]
def _get_psd_logm(mat, method):
    if isinstance(mat, torch.Tensor):
        eps = torch.tensor(torch.finfo(mat.dtype).eps)
        if method=='eigen':
            EVL,EVC = torch.linalg.eigh(mat)
            ret = (EVC * torch.log(torch.maximum(eps, EVL))) @ EVC.T.conj()
        else:
            assert (len(method)==3) and (method[0]=='pade')
            logm_op = numqi._torch_op.get_PSDMatrixLogm(int(method[1]), int(method[2]))
            ret = logm_op(mat)
    else: #numpy
        eps = np.finfo(mat.dtype).eps
        EVL,EVC = np.linalg.eigh(mat)
        ret = (EVC * np.log(np.maximum(eps, EVL))) @ EVC.T.conj()
    return ret
//  [python]
def FISTA(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver FISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            #L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()))[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

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

    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
        ahat_y = ahat_y.add(eta * basis.t().mm(Res))
        ahat = ahat_y.sub(eta * lambd).clamp(min = 0.)
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def test_orthogonal_iteration(n=100, d=30, k=10, bsz=10):
    torch.set_default_dtype(torch.float64)

    features = torch.randn(n, d)
    dataset = TensorDataset(features)
    loader = DataLoader(dataset, batch_size=bsz)

    eigenvalues, eigenvectors = numerical_distributed.orthogonal_iteration(
        input_mat=loader,
        k=k,
        num_power_iteration=400,
        dtype=torch.get_default_dtype(),
    )
    eigenvalues_expected, eigenvectors_expected = torch.symeig(features.T @ features)
    print(eigenvalues)
    print(eigenvalues_expected.flip(dims=(0,)))
    print('---')
    print(eigenvectors)
    print(eigenvectors_expected.flip(dims=(1,)))
    torch.testing.assert_allclose(eigenvalues, eigenvalues_expected.flip(dims=(0,))[:k], atol=1e-4, rtol=1e-4)
//  [python]
def test_orthogonal_iteration(n=100, d=30, k=10, bsz=10):
    torch.set_default_dtype(torch.float64)

    features = torch.randn(n, d)
    dataset = TensorDataset(features)
    loader = DataLoader(dataset, batch_size=bsz)

    eigenvalues, eigenvectors = numerical_distributed.orthogonal_iteration(
        input_mat=loader,
        k=k,
        num_power_iteration=400,
        dtype=torch.get_default_dtype(),
    )
    eigenvalues_expected, eigenvectors_expected = torch.linalg.eigh(features.T @ features)
    print(eigenvalues)
    print(eigenvalues_expected.flip(dims=(0,)))
    print('---')
    print(eigenvectors)
    print(eigenvectors_expected.flip(dims=(1,)))
    torch.testing.assert_allclose(eigenvalues, eigenvalues_expected.flip(dims=(0,))[:k], atol=1e-4, rtol=1e-4)
//  [python]
def forward(ctx, S_vec, triu_indices, d_indicies, dupl_mat):
        d = len(d_indicies)
        S = torch.zeros((*S_vec.shape[:-1], d, d))
        S[..., triu_indices[0], triu_indices[1]] = -S_vec
        S[..., triu_indices[1], triu_indices[0]] = S_vec
        D, V = torch.eig(1.0j * S, eigenvectors=True)  # j * S is Hermitian
        D = D[..., 0] / 1.0j
        exp_D = torch.exp(D)
        V_trns_conj = V.transpose(-2, -1).conj()
        ctx.save_for_backward(
            V, exp_D, V_trns_conj, D, triu_indices, d_indicies, dupl_mat
        )
        return (V * exp_D.unsqueeze(-2) @ V_trns_conj).real
//  [python]
def forward(ctx, S_vec, triu_indices, d_indicies, dupl_mat):
        d = len(d_indicies)
        S = torch.zeros((*S_vec.shape[:-1], d, d))
        S[..., triu_indices[0], triu_indices[1]] = -S_vec
        S[..., triu_indices[1], triu_indices[0]] = S_vec
        D, V = torch.linalg.eigh(1.0j * S)  # j * S is Hermitian
        D = D / 1.0j
        exp_D = torch.exp(D)
        V_trns_conj = V.transpose(-2, -1).conj()
        ctx.save_for_backward(
            V, exp_D, V_trns_conj, D, triu_indices, d_indicies, dupl_mat
        )
        return (V * exp_D.unsqueeze(-2) @ V_trns_conj).real
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False, decom_mode='svd') -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function          
        ensure_symmetric : bool = False (optional) 
            if ensure_symmetric=True, then M is symmetrized          
        ensure_psd : bool = False (optional) 
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0                  
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        if decom_mode=='eigh':
            s, U = torch.symeig(M,eigenvectors=True)
        elif decom_mode=='svd':
            U, s,_ = torch.linalg.svd(M)
        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False, decom_mode='svd') -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function          
        ensure_symmetric : bool = False (optional) 
            if ensure_symmetric=True, then M is symmetrized          
        ensure_psd : bool = False (optional) 
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0                  
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        if decom_mode=='eigh':
            s, U = torch.linalg.eigh(M)
        elif decom_mode=='svd':
            U, s,_ = torch.linalg.svd(M)
        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def create_frame(self, X, mask):
        mask = mask.unsqueeze(-1)
        center = (X * mask).sum(dim=1) / mask.sum(dim=1)
        X = X - center.unsqueeze(1) * mask  # [B, N, 3]
        C = torch.bmm(X.transpose(1, 2), X)  # [B, 3, 3] (Cov)
        L, V = torch.symeig(C.detach(), True)  # [B,3,3]
        # _, V = torch.linalg.eigh(C.float().detach(), UPLO='U')  # [B,3,3]
        F_ops = self.ops.unsqueeze(1).unsqueeze(0).to(X.device) * V.unsqueeze(1)  # [1,8,1,3] x [B,1,3,3] -> [B,8,3,3]
        h = torch.einsum('boij,bpj->bopi', F_ops.transpose(2, 3), X)  # transpose is inverse [B,8,N,3]
        h = h.view(X.size(0) * 8, X.size(1), 3)
        return h, F_ops.detach(), center
//  [python]
def create_frame(self, X, mask):
        mask = mask.unsqueeze(-1)
        center = (X * mask).sum(dim=1) / mask.sum(dim=1)
        X = X - center.unsqueeze(1) * mask  # [B, N, 3]
        C = torch.bmm(X.transpose(1, 2), X)  # [B, 3, 3] (Cov)
        L, V = torch.linalg.eigh(C.float().detach(), UPLO='U')  # [B,3,3]
        # _, V = torch.symeig(C.detach(), True)  # [B,3,3]
        F_ops = self.ops.unsqueeze(1).unsqueeze(0).to(X.device) * V.unsqueeze(1)  # [1,8,1,3] x [B,1,3,3] -> [B,8,3,3]
        h = torch.einsum('boij,bpj->bopi', F_ops.transpose(2, 3), X)  # transpose is inverse [B,8,N,3]
        h = h.view(X.size(0) * 8, X.size(1), 3)
        return h, F_ops.detach(), center
//  [python]
def get_fidelity(rho0:np.ndarray|torch.Tensor, rho1:np.ndarray|torch.Tensor):
    r'''get the fidelity of two density matrices or pure states

    Parameters:
        rho0 (np.ndarray,torch.Tensor): pure state (ndim=1) or density matrix (ndim=2)
        rho1 (np.ndarray,torch.Tensor): pure state or density matrix

    Returns:
        ret (float,torch.Tensor): the fidelity of the two states
    '''
    ndim0 = rho0.ndim
    ndim1 = rho1.ndim
    assert (ndim0 in {1,2}) and (ndim1 in {1,2})
    if isinstance(rho0, torch.Tensor):
        if ndim0==1 and ndim1==1:
            ret = torch.abs(torch.vdot(rho0, rho1))**2
        elif ndim0==1 and ndim1==2:
            ret = torch.vdot(rho0, rho1 @ rho0).real
        elif ndim0==2 and ndim1==1:
            ret = torch.vdot(rho1, rho0 @ rho1).real
        else:
            EVL0,EVC0 = torch.symeig(rho0)
            zero = torch.tensor(0.0, device=rho0.device)
            tmp0 = torch.sqrt(torch.maximum(zero, EVL0))
            tmp1 = (tmp0.reshape(-1,1) * EVC0.T.conj()) @ rho1 @ (EVC0 * tmp0)
            tmp2 = torch.symeig(tmp1).eigenvalues
            ret = torch.sum(torch.sqrt(torch.maximum(zero, tmp2)))**2
    else:
        if ndim0==1 and ndim1==1:
            ret = abs(np.vdot(rho0, rho1))**2
        elif ndim0==1 and ndim1==2:
            ret = np.vdot(rho0, rho1 @ rho0).real.item()
        elif ndim0==2 and ndim1==1:
            ret = np.vdot(rho1, rho0 @ rho1).real.item()
        else:
            EVL0,EVC0 = np.linalg.eigh(rho0)
            tmp0 = np.sqrt(np.maximum(0, EVL0))
            tmp1 = (tmp0[:,np.newaxis] * EVC0.T.conj()) @ rho1 @ (EVC0 * tmp0)
            tmp2 = np.linalg.eigvalsh(tmp1)
            ret = np.sum(np.sqrt(np.maximum(0, tmp2)))**2
    return ret
//  [python]
def get_fidelity(rho0:np.ndarray|torch.Tensor, rho1:np.ndarray|torch.Tensor):
    r'''get the fidelity of two density matrices or pure states

    Parameters:
        rho0 (np.ndarray,torch.Tensor): pure state (ndim=1) or density matrix (ndim=2)
        rho1 (np.ndarray,torch.Tensor): pure state or density matrix

    Returns:
        ret (float,torch.Tensor): the fidelity of the two states
    '''
    ndim0 = rho0.ndim
    ndim1 = rho1.ndim
    assert (ndim0 in {1,2}) and (ndim1 in {1,2})
    if isinstance(rho0, torch.Tensor):
        if ndim0==1 and ndim1==1:
            ret = torch.abs(torch.vdot(rho0, rho1))**2
        elif ndim0==1 and ndim1==2:
            ret = torch.vdot(rho0, rho1 @ rho0).real
        elif ndim0==2 and ndim1==1:
            ret = torch.vdot(rho1, rho0 @ rho1).real
        else:
            EVL0,EVC0 = torch.linalg.eigh(rho0)
            zero = torch.tensor(0.0, device=rho0.device)
            tmp0 = torch.sqrt(torch.maximum(zero, EVL0))
            tmp1 = (tmp0.reshape(-1,1) * EVC0.T.conj()) @ rho1 @ (EVC0 * tmp0)
            tmp2 = torch.linalg.eigvalsh(tmp1)
            ret = torch.sum(torch.sqrt(torch.maximum(zero, tmp2)))**2
    else:
        if ndim0==1 and ndim1==1:
            ret = abs(np.vdot(rho0, rho1))**2
        elif ndim0==1 and ndim1==2:
            ret = np.vdot(rho0, rho1 @ rho0).real.item()
        elif ndim0==2 and ndim1==1:
            ret = np.vdot(rho1, rho0 @ rho1).real.item()
        else:
            EVL0,EVC0 = np.linalg.eigh(rho0)
            tmp0 = np.sqrt(np.maximum(0, EVL0))
            tmp1 = (tmp0[:,np.newaxis] * EVC0.T.conj()) @ rho1 @ (EVC0 * tmp0)
            tmp2 = np.linalg.eigvalsh(tmp1)
            ret = np.sum(np.sqrt(np.maximum(0, tmp2)))**2
    return ret
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    return t.linalg.eigh(t_mat)
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    return torch.linalg.eigh(t_mat)
//  [python]
def obtain_attributes(data, use_adj=False, threshold=0.1, num_dim=32):
    save_node_border = 30000
        
    if use_adj:
        # to undirected and remove self-loop
        edges = to_undirected(data.edge_index)
        edges, _ = remove_self_loops(edges)
        tmp = to_dense_adj(edges)[0]
    else:
        tmp = similarity(data.x, data.x)
        
        # discretize the similarity matrix by threshold
        tmp = torch.where(tmp>threshold, 1.0, 0.0)

    tmp = get_laplacian_matrix(tmp)
    if tmp.shape[0] > save_node_border:
        L, V = scipy.linalg.eigh(tmp)
        L = torch.from_numpy(L)
        V = torch.from_numpy(V)
    else:
        L, V = torch.eig(tmp, eigenvectors=True) # much faster than torch.linalg.eig
    
    x = V[:, :num_dim].float()
    import sklearn.preprocessing as preprocessing
    x = preprocessing.normalize(x.cpu(), norm="l2")
    x = torch.tensor(x, dtype=torch.float32)

    return x
//  [python]
def obtain_attributes(data, use_adj=False, threshold=0.1, num_dim=32):
    save_node_border = 30000
        
    if use_adj:
        # to undirected and remove self-loop
        edges = to_undirected(data.edge_index)
        edges, _ = remove_self_loops(edges)
        tmp = to_dense_adj(edges)[0]
    else:
        tmp = similarity(data.x, data.x)
        
        # discretize the similarity matrix by threshold
        tmp = torch.where(tmp>threshold, 1.0, 0.0)

    tmp = get_laplacian_matrix(tmp)
    if tmp.shape[0] > save_node_border:
        L, V = scipy.linalg.eigh(tmp)
        L = torch.from_numpy(L)
        V = torch.from_numpy(V)
    else:
        L, V = torch.linalg.eigh(tmp) # much faster than torch.linalg.eig
    
    x = V[:, :num_dim].float()
    import sklearn.preprocessing as preprocessing
    x = preprocessing.normalize(x.cpu(), norm="l2")
    x = torch.tensor(x, dtype=torch.float32)

    return x
//  [python]
def get_relative_entropy(rho, sigma, tr_rho_log_rho:float=None, _torch_logm=('pade',6,8)):
    r'''get the relative entropy of two density matrices
    [wiki-link](https://en.wikipedia.org/wiki/Quantum_relative_entropy)

    $$ S(\rho,\sigma) = \mathrm{Tr}(\rho \log\rho - \rho \log\sigma) $$

    Parameters:
        rho (np.ndarray,torch.Tensor): a density matrix, shape=(dim,dim)
        sigma (np.ndarray,torch.Tensor): a density matrix, shape=(dim,dim)
        tr_rho_log_rho (float,None): tr(rho log(rho)), if None, calculate it
        _torch_logm (str,tuple): 'eigen' or ('pade',num_sqrtm,pade_order), 'pade' is used only when requires_grad

    Returns:
        ret (float,torch.Tensor): the relative entropy of the density matrices
    '''
    is_torch = isinstance(rho, torch.Tensor)
    if is_torch:
        eps = torch.tensor(torch.finfo(rho.dtype).eps)
        assert (_torch_logm=='eigen') or ((len(_torch_logm)==3) and (_torch_logm[0]=='pade'))
        if (rho.requires_grad or sigma.requires_grad) and (_torch_logm!='eigen'):
            tmp0 = numqi._torch_op.get_PSDMatrixLogm(int(_torch_logm[1]), int(_torch_logm[2]))
            log_sigma = tmp0(sigma)
        else:
            EVL,EVC = torch.eig(sigma,eigenvectors=True)
            EVL = EVL[:,0]
            log_sigma = (EVC * torch.log(torch.maximum(eps, EVL))) @ EVC.T.conj()
        ret = - torch.vdot(rho.reshape(-1), log_sigma.reshape(-1)).real
        if tr_rho_log_rho is None:
            EVL = torch.maximum(eps, torch.eigvalsh(rho))
            ret = ret + torch.dot(EVL, torch.log(EVL))
        else:
            ret = tr_rho_log_rho + ret
    else: #numpy
        eps = np.finfo(rho.dtype).eps
        EVL,EVC = np.linalg.eigh(sigma)
        log_sigma = (EVC * np.log(np.maximum(eps, EVL))) @ EVC.T.conj()
        ret = - np.vdot(rho.reshape(-1), log_sigma.reshape(-1)).real
        if tr_rho_log_rho is None:
            EVL = np.maximum(eps, np.linalg.eigvalsh(rho))
            ret = ret + np.dot(EVL, np.log(EVL))
        else:
            ret = tr_rho_log_rho + ret
    return ret
//  [python]
def get_relative_entropy(rho, sigma, tr_rho_log_rho:float=None, _torch_logm=('pade',6,8)):
    r'''get the relative entropy of two density matrices
    [wiki-link](https://en.wikipedia.org/wiki/Quantum_relative_entropy)

    $$ S(\rho,\sigma) = \mathrm{Tr}(\rho \log\rho - \rho \log\sigma) $$

    Parameters:
        rho (np.ndarray,torch.Tensor): a density matrix, shape=(dim,dim)
        sigma (np.ndarray,torch.Tensor): a density matrix, shape=(dim,dim)
        tr_rho_log_rho (float,None): tr(rho log(rho)), if None, calculate it
        _torch_logm (str,tuple): 'eigen' or ('pade',num_sqrtm,pade_order), 'pade' is used only when requires_grad

    Returns:
        ret (float,torch.Tensor): the relative entropy of the density matrices
    '''
    is_torch = isinstance(rho, torch.Tensor)
    if is_torch:
        eps = torch.tensor(torch.finfo(rho.dtype).eps)
        assert (_torch_logm=='eigen') or ((len(_torch_logm)==3) and (_torch_logm[0]=='pade'))
        if (rho.requires_grad or sigma.requires_grad) and (_torch_logm!='eigen'):
            tmp0 = numqi._torch_op.get_PSDMatrixLogm(int(_torch_logm[1]), int(_torch_logm[2]))
            log_sigma = tmp0(sigma)
        else:
            EVL,EVC = torch.linalg.eigh(sigma)
            log_sigma = (EVC * torch.log(torch.maximum(eps, EVL))) @ EVC.T.conj()
        ret = - torch.vdot(rho.reshape(-1), log_sigma.reshape(-1)).real
        if tr_rho_log_rho is None:
            EVL = torch.maximum(eps, torch.linalg.eigvalsh(rho))
            ret = ret + torch.dot(EVL, torch.log(EVL))
        else:
            ret = tr_rho_log_rho + ret
    else: #numpy
        eps = np.finfo(rho.dtype).eps
        EVL,EVC = np.linalg.eigh(sigma)
        log_sigma = (EVC * np.log(np.maximum(eps, EVL))) @ EVC.T.conj()
        ret = - np.vdot(rho.reshape(-1), log_sigma.reshape(-1)).real
        if tr_rho_log_rho is None:
            EVL = np.maximum(eps, np.linalg.eigvalsh(rho))
            ret = ret + np.dot(EVL, np.log(EVL))
        else:
            ret = tr_rho_log_rho + ret
    return ret
//  [python]
def compute_eig(self):
        assert self.cov is not None
        # TODO: Switch to a faster cuda eigh/svd implementation once available.
        # For now, use faster cpu eigh/svd implementation.
        # https://github.com/pytorch/pytorch/issues/41306
        # https://github.com/pytorch/pytorch/issues/69528
        # device = self.cov.device
        device = torch.device('cpu')
        cov = self.cov.to(device)
        try:
            eigvals, eigvecs = torch.symeig(cov, eigenvectors=True)
            # eigvecs, eigvals, _ = torch.linalg.svd(cov)
            # eigvecs = torch.flip(eigvecs, dims=[2])
            # eigvals = torch.flip(eigvals, dims=[1])
        except RuntimeError as ex:
            # https://github.com/pytorch/pytorch/issues/28293
            noise = 1e-6 * torch.rand_like(cov)
            eigvals, eigvecs = torch.symeig(cov + noise, eigenvectors=True)
        # Move back to original device.
        eigvals, eigvecs = eigvals.to(self.cov.device), eigvecs.to(self.cov.device)
        return eigvals, eigvecs
//  [python]
def compute_eig(self):
        assert self.cov is not None
        # TODO: Switch to a faster cuda eigh/svd implementation once available.
        # For now, use faster cpu eigh/svd implementation.
        # https://github.com/pytorch/pytorch/issues/41306
        # https://github.com/pytorch/pytorch/issues/69528
        # device = self.cov.device
        device = torch.device('cpu')
        cov = self.cov.to(device)
        try:
            eigvals, eigvecs = torch.linalg.eigh(cov)
            # eigvecs, eigvals, _ = torch.linalg.svd(cov)
            # eigvecs = torch.flip(eigvecs, dims=[2])
            # eigvals = torch.flip(eigvals, dims=[1])
        except RuntimeError as ex:
            # https://github.com/pytorch/pytorch/issues/28293
            noise = 1e-6 * torch.rand_like(cov)
            eigvals, eigvecs = torch.linalg.eigh(cov + noise)
        # Move back to original device.
        eigvals, eigvecs = eigvals.to(self.cov.device), eigvecs.to(self.cov.device)
        return eigvals, eigvecs
//  [python]
def test_eigh_svd_illcondition_matrix_input_should_not_crash(self, device, dtype):
        # See https://github.com/pytorch/pytorch/issues/94772, https://github.com/pytorch/pytorch/issues/105359
        # This test crashes with `cusolver error: CUSOLVER_STATUS_EXECUTION_FAILED` on cuda 11.8,
        # but passes on cuda 12.1 update 1 or later.
        a = torch.ones(512, 512, dtype=dtype, device=device)
        a[0, 0] = 1.0e-5
        a[-1, -1] = 1.0e5

        eigh_out = torch.linalg.eigh(a)
        svd_out = torch.linalg.svd(a)

        # Matrix input a is too ill-conditioned.
        # We'll just compare the first two singular values/eigenvalues. They are 1.0e5 and 511.0
        # The precision override with tolerance of 1.0 makes sense since ill-conditioned inputs are hard to converge
        # to exact values.
        self.assertEqual(eigh_out.eigenvalues.sort(descending=True).values[:2], [1.0e5, 511.0], atol=1.0, rtol=1.0e-2)
        self.assertEqual(svd_out.S[:2], [1.0e5, 511.0], atol=1.0, rtol=1.0e-2)
//  [python]
def patch_whitening(data, patch_size=(3, 3)):
    # Compute weights from data such that
    # torch.std(F.conv2d(data, weights), dim=(2, 3))
    # is close to 1.
    h, w = patch_size
    c = data.size(1)
    patches = data.unfold(2, h, 1).unfold(3, w, 1)
    patches = patches.transpose(1, 3).reshape(-1, c, h, w).to(torch.float32)

    n, c, h, w = patches.shape
    X = patches.reshape(n, c * h * w)
    X = X / (X.size(0) - 1) ** 0.5
    covariance = X.t() @ X

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

    eigenvalues = eigenvalues.flip(0)

    eigenvectors = eigenvectors.t().reshape(c * h * w, c, h, w).flip(0)

    return eigenvectors / torch.sqrt(eigenvalues + 1e-2).view(-1, 1, 1, 1)
//  [python]
def patch_whitening(data, patch_size=(3, 3)):
    # Compute weights from data such that
    # torch.std(F.conv2d(data, weights), dim=(2, 3))
    # is close to 1.
    h, w = patch_size
    c = data.size(1)
    patches = data.unfold(2, h, 1).unfold(3, w, 1)
    patches = patches.transpose(1, 3).reshape(-1, c, h, w).to(torch.float32)

    n, c, h, w = patches.shape
    X = patches.reshape(n, c * h * w)
    X = X / (X.size(0) - 1) ** 0.5
    covariance = X.t() @ X

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

    eigenvalues = eigenvalues.flip(0)

    eigenvectors = eigenvectors.t().reshape(c * h * w, c, h, w).flip(0)

    return eigenvectors / torch.sqrt(eigenvalues + 1e-2).view(-1, 1, 1, 1)
//  [python]
def torch_svd_eigh(matrix, rank=None):
	if rank is None: rank = min(matrix.shape[-2:])
	dim_1, dim_2 = matrix.shape[-2], matrix.shape[-1]
	kk = 1
	kkk = 1e-10
	if dim_2 > dim_1:
		X = matrix @ matrix.transpose(-1, -2)
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		S, U = torch.symeig(X, eigenvectors=True)
		S = S[..., -rank:].flip(-1).sub_(kk).clip_(min=kkk).sqrt_()
		U = U[..., -rank:].flip(-1)
		Vh = (U.transpose(-1, -2) @ matrix).div_(S[..., None])
	else:
		X = matrix.transpose(-1, -2) @ matrix
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		S, V = torch.symeig(X, eigenvectors=True)
		S = S[..., -rank:].flip(-1).sub_(kk).clip_(min=kkk).sqrt_()
		V = V[..., -rank:].flip(-1)
		U = (matrix @ V).div_(S[..., None, :])
		Vh = V.transpose(-1, -2)
	# assert torch.linalg.norm(U * S[..., None, :] @ Vh - matrix) / torch.linalg.norm(matrix) < 1e-6
	assert torch.linalg.norm(U.transpose(-1, -2) @ U - torch.eye(U.shape[-1], device=U.device)) / len(matrix) < 1e-6
	assert torch.linalg.norm(Vh @ Vh.transpose(-1, -2) - torch.eye(Vh.shape[-2], device=Vh.device)) / len(matrix) < 1e-6
	return U, S, Vh
//  [python]
def torch_svd_eigh(matrix, rank=None):
	if rank is None: rank = min(matrix.shape[-2:])
	dim_1, dim_2 = matrix.shape[-2], matrix.shape[-1]
	kk = 1
	kkk = 1e-10
	if dim_2 > dim_1:
		X = matrix @ matrix.transpose(-1, -2)
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		S, U = torch.linalg.eigh(X)
		S = S[..., -rank:].flip(-1).sub_(kk).clip_(min=kkk).sqrt_()
		U = U[..., -rank:].flip(-1)
		Vh = (U.transpose(-1, -2) @ matrix).div_(S[..., None])
	else:
		X = matrix.transpose(-1, -2) @ matrix
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		S, V = torch.linalg.eigh(X)
		S = S[..., -rank:].flip(-1).sub_(kk).clip_(min=kkk).sqrt_()
		V = V[..., -rank:].flip(-1)
		U = (matrix @ V).div_(S[..., None, :])
		Vh = V.transpose(-1, -2)
	# assert torch.linalg.norm(U * S[..., None, :] @ Vh - matrix) / torch.linalg.norm(matrix) < 1e-6
	assert torch.linalg.norm(U.transpose(-1, -2) @ U - torch.eye(U.shape[-1], device=U.device)) / len(matrix) < 1e-6
	assert torch.linalg.norm(Vh @ Vh.transpose(-1, -2) - torch.eye(Vh.shape[-2], device=Vh.device)) / len(matrix) < 1e-6
	return U, S, Vh
//  [python]
def make_positive_definite(
        x: Tensor,
        eps: Optional[float] = None) -> Tensor:
    """Makes the input positive definite."""
    if eps is None:
        eps = torch.finfo(x.dtype).eps
    try:
        eigval, eigvec = torch.symeig(x)
        eigval = torch.maximum(eigval, torch.as_tensor(eps).to(eigval))
        return eigvec @ eigval.diag_embed() @ eigvec.t()
    except RuntimeError:
        print("PyTorch eigh does not converge. Adding eps to diagonal.")
        return make_positive_definite2(x, eps)
//  [python]
def make_positive_definite(
        x: Tensor,
        eps: Optional[float] = None) -> Tensor:
    """Makes the input positive definite."""
    if eps is None:
        eps = torch.finfo(x.dtype).eps
    try:
        eigval, eigvec = torch.linalg.eigh(x)
        eigval = torch.maximum(eigval, torch.as_tensor(eps).to(eigval))
        return eigvec @ eigval.diag_embed() @ eigvec.transpose(-2, -1)
    except RuntimeError:
        print("PyTorch eigh does not converge. Adding eps to diagonal.")
        return make_positive_definite2(x, eps)
//  [python]
def PCA_eig(X, k: int, scale: bool = False):
    m, t, p = X.size()
    X = X.reshape(m * t, p)
    n, p = X.size()
    X_center = X - X.mean(1).view(n, 1)
    covariance = 1 / (n - 1) * torch.mm(X_center.t(), X_center)
    if scale:
        covariance = torch.mm(torch.diag(torch.sqrt(1 / torch.diag(covariance))), covariance)
    eigenvalues, eigenvectors = torch.symeig(covariance, eigenvectors=True)
    components = eigenvectors[:, :k]
    X = X @ components
    return X.reshape(m, t, k)
//  [python]
def PCA_eig(X, k: int, scale: bool = False):
    m, t, p = X.size()
    X = X.reshape(m * t, p)
    n, p = X.size()
    X_center = X - X.mean(1).view(n, 1)
    covariance = 1 / (n - 1) * torch.mm(X_center.t(), X_center)
    if scale:
        covariance = torch.mm(torch.diag(torch.sqrt(1 / torch.diag(covariance))), covariance)
    _, eigenvectors = torch.linalg.eigh(covariance)  # cov is symmetric => eigh OK
    components = eigenvectors[:, :k]
    X = X @ components
    return X.reshape(m, t, k)
//  [python]
def _calculate_scaled_laplacian_tensor(self, adj, lambda_max=2, undirected=True):
        if undirected:
            adj = torch.maximum(adj, adj.T)
        L = self._calculate_laplacian_tensor(adj)

        # calculate the maximum eigenvalue of adjacent matrix
        if lambda_max is None:
            eigenvalues, _ = torch.symeig(L)
            lambda_max = eigenvalues.max()
        
        I = torch.eye(L.shape[0], device=self._device)
        L = (2 / lambda_max * L ) - I 
        return L
//  [python]
def _calculate_scaled_laplacian_tensor(self, adj, lambda_max=2, undirected=True):
        if undirected:
            adj = torch.maximum(adj, adj.T)
        L = self._calculate_laplacian_tensor(adj)

        # calculate the maximum eigenvalue of adjacent matrix
        if lambda_max is None:
            eigenvalues, _ = torch.linalg.eigh(L)
            lambda_max = eigenvalues.max()
        
        I = torch.eye(L.shape[0], device=self._device)
        L = (2 / lambda_max * L ) - I 
        return L
//  [python]
def EVD_Laplacian(data, norm=None):
    L_raw = get_laplacian(to_undirected(data.edge_index, num_nodes=data.num_nodes), 
                          normalization=norm, num_nodes=data.num_nodes)
    L = SparseTensor(row=L_raw[0][0], col=L_raw[0][1], value=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    # L_raw = to_undirected(data.edge_index, num_nodes=data.num_nodes)
    # L = SparseTensor(row=L_raw[0], col=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    D, V  = torch.symeig(L, eigenvectors=True)
    return D, V
//  [python]
def EVD_Laplacian(data, norm=None):
    L_raw = get_laplacian(to_undirected(data.edge_index, num_nodes=data.num_nodes), 
                          normalization=norm, num_nodes=data.num_nodes)
    L = SparseTensor(row=L_raw[0][0], col=L_raw[0][1], value=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    # L_raw = to_undirected(data.edge_index, num_nodes=data.num_nodes)
    # L = SparseTensor(row=L_raw[0], col=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    D, V  = torch.linalg.eigh(L)
    return D, V
//  [python]
def G_WF_L(H,F,deriv=False):
    """
    Computes some derived model parameters.
    
    Inputs (model parameters):
        H: (D,D), used in W = H @ H.T, where W is within-class precision
        F: (D,d)  D > d, speaker factor loading matrix
        
    Outputs (derived model parameters):
        
        G: (D,D): To be used in x'Gx, to project x-vector energy (squared 
                  magnitude) into complement of speaker subspace. This is used 
                  elsewhere to compute b, the precision scaling factor.
        WF: (D,d) To project x-vector into speaker subspace. This is used
                  elsewehere to compute the first-order stats A.          
                  
        Eigen: a tuple containing the eigenanalysis of B
                  
    """
    W = H @ H.T            # (D,D)
    WF = W @ F             # (D,d)
    B = F.T @ WF           # (d,d) 
    E, V = torch.symeig(B, eigenvectors=True) # B = V @ E @ V.T
    Ecol = E.reshape(-1,1)
    
    Eigen = E, V, Ecol 
    
    def solveB(RHS):
        return V @ ((V.T @ RHS)/Ecol)
        
    S = solveB(WF.T) 
    G = W - WF @ S      
    L = F.T @ H
    if not deriv: return G, WF, Eigen
    def back(dG,dWF,dL):
    
        # L = F.T @ H
        dL = dL @ L              # perform deferred multiplication
        dF = H @ dL.T
        dH = F @ dL
        
        # G = W - WF @ S
        dW = dG
        dWF = dWF - dG @ S.T    # avoiding side-effect on dWF
        dS = -WF.T @ dG
        
        #S = solveB(WF.T)
        dRHS = solveB(dS)     # B is symmetric
        dWF += dRHS.T
        dB = -dRHS @ S.T
        
        # B = F.T @ WF    ==> B.T = WF.T @ F 
        dWF += F @ dB
        dF += WF @ dB.T
        
        # WF = W @ F
        dW = dW + dWF @ F.T       # avoiding side-effect on dG
        dF += W.T @ dWF
        
        # W = H @ H.T
        dH += (dW+dW.T) @ H
        return dH, dF
    return G, WF, Eigen, back
//  [python]
def G_WF_L(H,F,deriv=False):
    """
    Computes some derived model parameters.
    
    Inputs (model parameters):
        H: (D,D), used in W = H @ H.T, where W is within-class precision
        F: (D,d)  D > d, speaker factor loading matrix
        
    Outputs (derived model parameters):
        
        G: (D,D): To be used in x'Gx, to project x-vector energy (squared 
                  magnitude) into complement of speaker subspace. This is used 
                  elsewhere to compute b, the precision scaling factor.
        WF: (D,d) To project x-vector into speaker subspace. This is used
                  elsewehere to compute the first-order stats A.          
                  
        Eigen: a tuple containing the eigenanalysis of B
                  
    """
    W = H @ H.T            # (D,D)
    WF = W @ F             # (D,d)
    B = F.T @ WF           # (d,d) 
    E, V = torch.linalg.eigh(B, UPLO='U') # B = V @ E @ V.T
    Ecol = E.reshape(-1,1)
    
    Eigen = E, V, Ecol 
    
    def solveB(RHS):
        return V @ ((V.T @ RHS)/Ecol)
        
    S = solveB(WF.T) 
    G = W - WF @ S      
    L = F.T @ H
    if not deriv: return G, WF, Eigen
    def back(dG,dWF,dL):
    
        # L = F.T @ H
        dL = dL @ L              # perform deferred multiplication
        dF = H @ dL.T
        dH = F @ dL
        
        # G = W - WF @ S
        dW = dG
        dWF = dWF - dG @ S.T    # avoiding side-effect on dWF
        dS = -WF.T @ dG
        
        #S = solveB(WF.T)
        dRHS = solveB(dS)     # B is symmetric
        dWF += dRHS.T
        dB = -dRHS @ S.T
        
        # B = F.T @ WF    ==> B.T = WF.T @ F 
        dWF += F @ dB
        dF += WF @ dB.T
        
        # WF = W @ F
        dW = dW + dWF @ F.T       # avoiding side-effect on dG
        dF += W.T @ dWF
        
        # W = H @ H.T
        dH += (dW+dW.T) @ H
        return dH, dF
    return G, WF, Eigen, back
//  [python]
def get_pos_enc(self, idx):
        if idx in self.pos_enc:
            return self.pos_enc[idx]
        else:
            graph_info = self.graphs[idx]
            edge_index = graph_info['edge_index']
            n_atoms = graph_info['num_nodes']
            sparse = torch.sparse_coo_tensor(edge_index, torch.ones(edge_index.shape[1]), size=(n_atoms, n_atoms),
                                             device=self.device)
            A = sparse.to_dense()
            A += torch.eye(n_atoms, device=self.device)
            D = A.sum(dim=0)
            D_mat = torch.diag(D)
            L = D_mat - A
            D += (D == 0)
            N = D ** -0.5
            L_sym = torch.eye(n_atoms, device=self.device) - N * L * N

            eig_vals, eig_vecs = torch.symeig(L_sym, eigenvectors=True)
            idx = eig_vals.argsort()[0: self.num_freq]  # Keep up to the maximum desired number of frequencies
            eig_vals, eig_vecs = eig_vals[idx], eig_vecs[:, idx]

            # Sort, normalize and pad EigenVectors
            eig_vecs = eig_vecs[:, eig_vals.argsort()]  # increasing order
            eig_vecs = F.normalize(eig_vecs, p=2, dim=1, eps=1e-12, out=None)
            if n_atoms < self.num_freq:
                eig_vecs = F.pad(eig_vecs, (0, self.num_freq - n_atoms), value=float('nan'))
                eig_vals = F.pad(eig_vals, (0, self.num_freq - n_atoms), value=float('nan'))
            self.pos_enc[idx] = (eig_vals.cpu(), eig_vecs.cpu())
            return self.pos_enc[idx]
//  [python]
def get_pos_enc(self, idx):
        if idx in self.pos_enc:
            return self.pos_enc[idx]
        else:
            graph_info = self.graphs[idx]
            edge_index = graph_info['edge_index']
            n_atoms = graph_info['num_nodes']
            sparse = torch.sparse_coo_tensor(edge_index, torch.ones(edge_index.shape[1]), size=(n_atoms, n_atoms),
                                             device=self.device)
            A = sparse.to_dense()
            A += torch.eye(n_atoms, device=self.device)
            D = A.sum(dim=0)
            D_mat = torch.diag(D)
            L = D_mat - A
            D += (D == 0)
            N = D ** -0.5
            L_sym = torch.eye(n_atoms, device=self.device) - N * L * N

            eig_vals, eig_vecs = torch.linalg.eigh(L_sym)
            idx = eig_vals.argsort()[0: self.num_freq]  # Keep up to the maximum desired number of frequencies
            eig_vals, eig_vecs = eig_vals[idx], eig_vecs[:, idx]

            # Sort, normalize and pad EigenVectors
            eig_vecs = eig_vecs[:, eig_vals.argsort()]  # increasing order
            eig_vecs = F.normalize(eig_vecs, p=2, dim=1, eps=1e-12, out=None)
            if n_atoms < self.num_freq:
                eig_vecs = F.pad(eig_vecs, (0, self.num_freq - n_atoms), value=float('nan'))
                eig_vals = F.pad(eig_vals, (0, self.num_freq - n_atoms), value=float('nan'))
            self.pos_enc[idx] = (eig_vals.cpu(), eig_vecs.cpu())
            return self.pos_enc[idx]
//  [python]
def ss_update(self,SExx,n,lr=1.0):
        idx = n>1
        SExx = SExx*(idx).unsqueeze(-1).unsqueeze(-1)
        invU = (self.invU_0 + SExx)*lr + (1-lr)*self.invU
        self.nu = (self.nu_0 + n)*lr + (1-lr)*self.nu
        self.d, self.v = torch.symeig(0.5*invU+0.5*invU.transpose(-2,-1))
//  [python]
def ss_update(self,SExx,n,lr=1.0):
        idx = n>1
        SExx = SExx*(idx).unsqueeze(-1).unsqueeze(-1)
        invU = (self.invU_0 + SExx)*lr + (1-lr)*self.invU
        self.nu = (self.nu_0 + n)*lr + (1-lr)*self.nu
        self.d, self.v = torch.linalg.eigh(0.5*invU+0.5*invU.transpose(-2,-1))
//  [python]
def scatter_pca(x, idx, on_cpu=True):
    """Scatter implementation for PCA.

    Returns eigenvalues and eigenvectors for each group in idx.
    If x has shape N1xD and idx covers indices in [0, N2], the
    eigenvalues will have shape N2xD and the eigenvectors will
    have shape N2xDxD. The eigenvalues and eigenvectors are
    sorted by increasing eigenvalue.
    """
    assert idx.dim() == 1
    assert x.dim() == 2
    assert idx.shape[0] == x.shape[0]
    assert x.shape[1] > 1

    d = x.shape[1]
    device = x.device

    # Substract mean
    mean = scatter_mean(x, idx, dim=0)
    x = x - mean[idx]

    # Compute pointwise covariance as a N_1x(DxD) matrix
    ij = torch.tensor(list(combinations_with_replacement(range(d), 2)), device=device)
    upper_triangle = x[:, ij[:, 0]] * x[:, ij[:, 1]]

    # Aggregate the covariances as a N_2x(DxD) with scatter_sum
    # and convert it to a N_2xDxD batch of matrices
    upper_triangle = scatter_add(upper_triangle, idx, dim=0) / d
    cov = torch.empty((upper_triangle.shape[0], d, d), device=device)
    cov[:, ij[:, 0], ij[:, 1]] = upper_triangle

    # Eigendecompostion
    if on_cpu:
        device = cov.device
        cov = cov.cpu()
        eval, evec = torch.symeig(cov, eigenvectors=True)
        eval = eval.to(device)
        evec = evec.to(device)
    else:
        eval, evec = torch.symeig(cov, eigenvectors=True)

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

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

    return eval, evec
//  [python]
def scatter_pca(x, idx, on_cpu=True):
    """Scatter implementation for PCA.

    Returns eigenvalues and eigenvectors for each group in idx.
    If x has shape N1xD and idx covers indices in [0, N2], the
    eigenvalues will have shape N2xD and the eigenvectors will
    have shape N2xDxD. The eigenvalues and eigenvectors are
    sorted by increasing eigenvalue.
    """
    assert idx.dim() == 1
    assert x.dim() == 2
    assert idx.shape[0] == x.shape[0]
    assert x.shape[1] > 1

    d = x.shape[1]
    device = x.device

    # Substract mean
    mean = scatter_mean(x, idx, dim=0)
    x = x - mean[idx]

    # Compute pointwise covariance as a N_1x(DxD) matrix
    ij = torch.tensor(list(combinations_with_replacement(range(d), 2)), device=device)
    upper_triangle = x[:, ij[:, 0]] * x[:, ij[:, 1]]

    # Aggregate the covariances as a N_2x(DxD) with scatter_sum
    # and convert it to a N_2xDxD batch of matrices
    upper_triangle = scatter_add(upper_triangle, idx, dim=0) / d
    cov = torch.empty((upper_triangle.shape[0], d, d), device=device)
    cov[:, ij[:, 0], ij[:, 1]] = upper_triangle

    # Eigendecompostion
    if on_cpu:
        device = cov.device
        cov = cov.cpu()
        eval, evec = torch.linalg.eigh(cov, UPLO='U')
        eval = eval.to(device)
        evec = evec.to(device)
    else:
        eval, evec = torch.linalg.eigh(cov, UPLO='U')

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

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

    return eval, evec
//  [python]
def __init__(self,nu,U):  #nu, invU are natural parameters, nu*U is expected value of precision matrix

        self.dim = U.shape[-1]
        self.event_dim = 2
        self.batch_dim = U.ndim-2
        self.event_shape = U.shape[-2:]
        self.batch_shape = U.shape[:-2]

        self.d, self.v = torch.symeig(U,eigenvectors=True)
        self.d = 1.0/self.d
        self.invU_0 = self.v@(self.d.unsqueeze(-1)*self.v.transpose(-2,-1))
        self.logdet_invU_0 = self.d.log().sum(-1)
        self.nu_0 = nu
        self.nu = self.nu_0
//  [python]
def __init__(self,nu,U):  #nu, invU are natural parameters, nu*U is expected value of precision matrix

        self.dim = U.shape[-1]
        self.event_dim = 2
        self.batch_dim = U.ndim-2
        self.event_shape = U.shape[-2:]
        self.batch_shape = U.shape[:-2]

        self.d, self.v = torch.linalg.eigh(U)
        self.d = 1.0/self.d
        self.invU_0 = self.v@(self.d.unsqueeze(-1)*self.v.transpose(-2,-1))
        self.logdet_invU_0 = self.d.log().sum(-1)
        self.nu_0 = nu
        self.nu = self.nu_0
//  [python]
def set_inv_mass_matrix(self, mat, batch_dim=False):
        """ Sets the inverse mass matrix and defines corresponding kinetic energy function and its gradient.

        Args:
            mat (torch.Tensor or None): Diagonal elements or full matrix of the inverse mass matrix.
            batch_dim (bool, optional): Is there a batch dimension? Defaults to False.
        """
        if mat is None:
            self.mass_matrix_inv = None
            self.mass_matrix_sqrt = None

            # Define kinetic energy and its gradient
            self.KE = lambda p: 0.5*(p**2).sum(-1) # Sum across parameter dimension
            self.KE_g = lambda p: p
        else:
            assert mat.ndim == 1 + int(batch_dim) or mat.ndim == 2 + int(batch_dim)
            if batch_dim:
                self.mass_matrix_batch_dim = True
            if mat.ndim == 1 + int(batch_dim):
                self.mass_matrix_inv = mat
                self.mass_matrix_sqrt = torch.stack([torch.diag(m ** -0.5) for m in mat], dim=0)

                # Define kinetic energy and its gradient
                self.KE = lambda p: 0.5*(p**2 * self.mass_matrix_inv).sum(-1) # Sum across parameter dimension
                self.KE_g = lambda p: p * self.mass_matrix_inv
            else:
                L_list, Q_list = [], []
                mass_matrix_sqrt_list = []

                for i,m in enumerate(mat):
                    assert (m == m.T).all() # Check if symmetric
                    L, Q = torch.symeig(m)
                    if (L > 0).all(): # Check if positive definite
                        L_list.append(L)
                        Q_list.append(Q)
                        mass_matrix_sqrt_list.append(Q @ torch.diag(L**-0.5) @ Q.T)
                    else:
                        m = mat.mean(dim=0)
                        L, Q = torch.symeig(m)
                        L_list.append(L)
                        Q_list.append(Q)
                        mass_matrix_sqrt_list.append(Q @ torch.diag(L**-0.5) @ Q.T)
                        mat[i] = m
                
                self.mass_matrix_inv = mat
                self.mass_matrix_sqrt = torch.stack(mass_matrix_sqrt_list, dim=0)

                # Define kinetic energy and its gradient
                self.KE = lambda p: 0.5*(p * (self.mass_matrix_inv @ p.unsqueeze(-1)).squeeze(-1)).sum(-1) # Sum across parameter dimension
                self.KE_g = lambda p: (self.mass_matrix_inv @ p.unsqueeze(-1)).squeeze(-1)
//  [python]
def set_inv_mass_matrix(self, mat, batch_dim=False):
        """ Sets the inverse mass matrix and defines corresponding kinetic energy function and its gradient.

        Args:
            mat (torch.Tensor or None): Diagonal elements or full matrix of the inverse mass matrix.
            batch_dim (bool, optional): Is there a batch dimension? Defaults to False.
        """
        if mat is None:
            self.mass_matrix_inv = None
            self.mass_matrix_sqrt = None

            # Define kinetic energy and its gradient
            self.KE = lambda p: 0.5*(p**2).sum(-1) # Sum across parameter dimension
            self.KE_g = lambda p: p
        else:
            assert mat.ndim == 1 + int(batch_dim) or mat.ndim == 2 + int(batch_dim)
            if batch_dim:
                self.mass_matrix_batch_dim = True
            if mat.ndim == 1 + int(batch_dim):
                self.mass_matrix_inv = mat
                self.mass_matrix_sqrt = torch.stack([torch.diag(m ** -0.5) for m in mat], dim=0)

                # Define kinetic energy and its gradient
                self.KE = lambda p: 0.5*(p**2 * self.mass_matrix_inv).sum(-1) # Sum across parameter dimension
                self.KE_g = lambda p: p * self.mass_matrix_inv
            else:
                L_list, Q_list = [], []
                mass_matrix_sqrt_list = []

                for i,m in enumerate(mat):
                    assert (m == m.T).all() # Check if symmetric
                    L, Q = torch.linalg.eigh(m)
                    if (L > 0).all(): # Check if positive definite
                        L_list.append(L)
                        Q_list.append(Q)
                        mass_matrix_sqrt_list.append(Q @ torch.diag(L**-0.5) @ Q.T)
                    else:
                        m = mat.mean(dim=0)
                        L, Q = torch.linalg.eigh(m)
                        L_list.append(L)
                        Q_list.append(Q)
                        mass_matrix_sqrt_list.append(Q @ torch.diag(L**-0.5) @ Q.T)
                        mat[i] = m
                
                self.mass_matrix_inv = mat
                self.mass_matrix_sqrt = torch.stack(mass_matrix_sqrt_list, dim=0)

                # Define kinetic energy and its gradient
                self.KE = lambda p: 0.5*(p * (self.mass_matrix_inv @ p.unsqueeze(-1)).squeeze(-1)).sum(-1) # Sum across parameter dimension
                self.KE_g = lambda p: (self.mass_matrix_inv @ p.unsqueeze(-1)).squeeze(-1)
//  [python]
def nat_update(self,nu,invU):
        self.nu = nu
        self.d, self.v = torch.symeig(0.5*invU+0.5*invU.transpose(-2,-1))
//  [python]
def nat_update(self,nu,invU):
        self.nu = nu
        self.d, self.v = torch.linalg.eigh(0.5*invU+0.5*invU.transpose(-2,-1))