jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        # error happens
        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.

                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d' or classname == 'ConvTranspose2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            if self.Q_a[m].device != self.Q_g[m].device:
                self.Q_a[m]=self.Q_a[m].to(self.Q_g[m].device)
                self.d_a[m]=self.d_a[m].to(self.Q_g[m].device)

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v


        vg_sum = 0

        for p in self.model.parameters():
            # if len(p.size())>=3:
            #     continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()


        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            # if len(p.size())>=3:
            #     continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        # error happens
        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.

                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
                    self.m_gg[m])
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(
                    self.m_aa[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d' or classname == 'ConvTranspose2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            if self.Q_a[m].device != self.Q_g[m].device:
                self.Q_a[m]=self.Q_a[m].to(self.Q_g[m].device)
                self.d_a[m]=self.d_a[m].to(self.Q_g[m].device)

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v


        vg_sum = 0

        for p in self.model.parameters():
            # if len(p.size())>=3:
            #     continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()


        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            # if len(p.size())>=3:
            #     continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def eigh(
    x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    result_tuple = NamedTuple(
        "eigh", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)]
    )
    eigenvalues, eigenvectors = torch.symeig(x, eigenvectors=True, upper=UPLO == "U", out=out)
    return result_tuple(eigenvalues, eigenvectors)
//  [python]
def eigh(
    x: torch.Tensor, /, *, UPLO: str = "L", out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    result_tuple = NamedTuple(
        "eigh", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)]
    )
    eigenvalues, eigenvectors = torch.linalg.eigh(x, UPLO=UPLO, out=out)
    return result_tuple(eigenvalues, eigenvectors)
//  [python]
def loss_function(pastdata=None, futuredata=None, criterion='SRVNet', epsilon=1e-6, method='cutoff'):

    """
    Loss function calculation for the VAMPNet / SRVNet training.

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    criterion: string, default='SRVNet'
        The criterion used to calculate the score/loss fucntion, can be chosen from: 'SRVNet', 'VAMP1', 'VAMP2', 'VAMPE'
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    loss function: float
        The loss function for the set criterion.
    """

    _criterion = ['SRVNet', 'VAMP1', 'VAMP2', 'VAMPE']
    assert criterion in _criterion, "Invalid criterion "+str(criterion)+" was adopted, should use criterion in: "+str(_criterion)
    
    if criterion == 'SRVNet':
        rev_propogator = reverse_propogator(pastdata=pastdata, futuredata=futuredata, epsilon=0, method='regularize')
        _eigenval, _ = torch.symeig(rev_propogator)
        _score = torch.sum(torch.pow(_eigenval, 2)) + 1


    elif criterion == 'VAMP1':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method)
        _score = torch.norm(koopman_mat, p='nuc') + 1


    elif criterion == 'VAMP2':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method)
        _score = torch.pow(torch.norm(koopman_mat, p='fro'), 2) + 1


    elif criterion == 'VAMPE':
        c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True)
        c00_sqrt_inv = matrix_inverse(matrix=c00, sqrt_root=True, epsilon=epsilon, method=method)
        ctt_sqrt_inv = matrix_inverse(matrix=ctt, sqrt_root=True,  epsilon=epsilon, method=method)

        koopman_mat = torch.matmul(c00_sqrt_inv, torch.matmul(c0t, ctt_sqrt_inv)).t()
        u, sigma, v = torch.svd(koopman_mat)
        mask = sigma > epsilon
        
        u = torch.mm(c00_sqrt_inv, u[:, mask])
        v = torch.mm(ctt_sqrt_inv, v[:, mask])
        sigma = sigma[mask]

        u_t = u.t(); v_t = v.t()
        sigma = torch.diag(sigma)
        _score = torch.trace(2. * torch.linalg.multi_dot([sigma, u_t, c0t, v]) - torch.linalg.multi_dot([sigma, u_t, c00, u, sigma, v_t, ctt, v])) + 1

    return -1. * _score
//  [python]
def loss_function(pastdata=None, futuredata=None, criterion='SRVNet', epsilon=1e-6, method='cutoff'):

    """
    Loss function calculation for the VAMPNet / SRVNet training.

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    criterion: string, default='SRVNet'
        The criterion used to calculate the score/loss fucntion, can be chosen from: 'SRVNet', 'VAMP1', 'VAMP2', 'VAMPE'
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    loss function: float
        The loss function for the set criterion.
    """

    _criterion = ['SRVNet', 'VAMP1', 'VAMP2', 'VAMPE']
    assert criterion in _criterion, "Invalid criterion "+str(criterion)+" was adopted, should use criterion in: "+str(_criterion)
    
    if criterion == 'SRVNet':
        rev_propogator = reverse_propogator(pastdata=pastdata, futuredata=futuredata, epsilon=0, method='regularize')
        _eigenval, _ = torch.linalg.eigh(rev_propogator)
        _score = torch.sum(torch.pow(_eigenval, 2)) + 1


    elif criterion == 'VAMP1':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method)
        _score = torch.norm(koopman_mat, p='nuc') + 1


    elif criterion == 'VAMP2':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method)
        _score = torch.pow(torch.norm(koopman_mat, p='fro'), 2) + 1


    elif criterion == 'VAMPE':
        c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True)
        c00_sqrt_inv = matrix_inverse(matrix=c00, sqrt_root=True, epsilon=epsilon, method=method)
        ctt_sqrt_inv = matrix_inverse(matrix=ctt, sqrt_root=True,  epsilon=epsilon, method=method)

        koopman_mat = torch.matmul(c00_sqrt_inv, torch.matmul(c0t, ctt_sqrt_inv)).t()
        u, sigma, v = torch.svd(koopman_mat)
        mask = sigma > epsilon
        
        u = torch.mm(c00_sqrt_inv, u[:, mask])
        v = torch.mm(ctt_sqrt_inv, v[:, mask])
        sigma = sigma[mask]

        u_t = u.t(); v_t = v.t()
        sigma = torch.diag(sigma)
        _score = torch.trace(2. * torch.linalg.multi_dot([sigma, u_t, c0t, v]) - torch.linalg.multi_dot([sigma, u_t, c00, u, sigma, v_t, ctt, v])) + 1

    return -1. * _score
//  [python]
def _torch_psd_sqrtm_forward_repeat(matA, repeat:int=1):
    assert repeat >= 1
    shape = matA.shape
    assert shape[-1] == shape[-2]
    matA = matA.reshape(-1, shape[-1], shape[-1])
    # scipy.linalg.sqrtm is much slower than eigh
    # ret = matA.numpy()
    # if matA.ndim==2:
    #     for _ in range(repeat):
    #         # TODO .astype(matA.dtype) scipy-v1.10 bug https://github.com/scipy/scipy/issues/18250
    #         ret = scipy.linalg.sqrtm(ret).astype(ret.dtype)
    # ctx_tensor = ret,
    EVL, EVC = torch.symeig(matA)
    sqrt_EVL = torch.maximum(torch.zeros(1, dtype=EVL.dtype, device=EVL.device), EVL)
    for _ in range(repeat):
        sqrt_EVL = torch.sqrt(sqrt_EVL)
    ret = ((EVC*sqrt_EVL.view(-1,1,shape[-1])) @ EVC.transpose(1,2).conj()).view(*shape)
    ctx_tensor = sqrt_EVL, EVC
    return ret,ctx_tensor
//  [python]
def _torch_psd_sqrtm_forward_repeat(matA, repeat:int=1):
    assert repeat >= 1
    shape = matA.shape
    assert shape[-1] == shape[-2]
    matA = matA.reshape(-1, shape[-1], shape[-1])
    # scipy.linalg.sqrtm is much slower than eigh
    # ret = matA.numpy()
    # if matA.ndim==2:
    #     for _ in range(repeat):
    #         # TODO .astype(matA.dtype) scipy-v1.10 bug https://github.com/scipy/scipy/issues/18250
    #         ret = scipy.linalg.sqrtm(ret).astype(ret.dtype)
    # ctx_tensor = ret,
    EVL, EVC = torch.linalg.eigh(matA)
    sqrt_EVL = torch.maximum(torch.zeros(1, dtype=EVL.dtype, device=EVL.device), EVL)
    for _ in range(repeat):
        sqrt_EVL = torch.sqrt(sqrt_EVL)
    ret = ((EVC*sqrt_EVL.view(-1,1,shape[-1])) @ EVC.transpose(1,2).conj()).view(*shape)
    ctx_tensor = sqrt_EVL, EVC
    return ret,ctx_tensor
//  [python]
def read_graph_dgl(raw_dir, add_inverse_edge = False, additional_node_files = [], additional_edge_files = [], binary=False):

    if binary:
        # npz
        graph_list = read_binary_graph_raw(raw_dir, add_inverse_edge)
    else:
        # csv
        graph_list = read_csv_graph_raw(raw_dir, add_inverse_edge, additional_node_files = additional_node_files, additional_edge_files = additional_edge_files)
        
    dgl_graph_list = []

    print('Converting graphs into DGL objects...')
    
    for graph in tqdm(graph_list):

        src, dst = torch.from_numpy(graph['edge_index'])
        num_nodes = graph['num_nodes']

        if num_nodes == 1:  # some graphs have one node
            A_ = torch.tensor(1.).view(1, 1)
        else:
            A = torch.zeros([num_nodes, num_nodes], dtype=torch.float)
            A[src, dst] = 1.0
            for i in range(num_nodes):
                A[i, i] = 1.0
            deg = torch.sum(A, axis=0).squeeze() ** -0.5
            D = torch.diag(deg)
            A_ = D @ A @ D
        e, u = torch.symeig(A_, eigenvectors=True)

        fully_connected = torch.ones([num_nodes, num_nodes], dtype=torch.float).nonzero(as_tuple=True)
        g = dgl.graph(fully_connected, num_nodes = num_nodes)

        g.ndata['e'] = e
        g.ndata['u'] = u

        if graph['node_feat'] is not None:
            g.ndata['feat'] = torch.from_numpy(graph['node_feat'])

        if graph['edge_feat'] is not None:
            edge_idx = torch.stack([src, dst], dim=0)
            edge_attr = torch.from_numpy(graph['edge_feat']) + 1  # for padding

            if len(edge_attr.shape) == 1:
                edge_attr_dense = to_dense_adj(edge_idx, edge_attr=edge_attr.unsqueeze(-1)).squeeze(0).squeeze(-1).view(-1)
            else:
                if edge_attr.size(0) == 0:
                    edge_attr_dense = torch.zeros([num_nodes ** 2, edge_attr.size(1)]).long()  # for graphs without edge
                else:
                    edge_attr_dense = to_dense_adj(edge_idx, edge_attr=edge_attr, max_num_nodes=num_nodes).squeeze(0).view(-1, edge_attr.shape[-1])

            g.edata['feat'] = edge_attr_dense

        dgl_graph_list.append(g)

    return dgl_graph_list
//  [python]
def read_graph_dgl(raw_dir, add_inverse_edge = False, additional_node_files = [], additional_edge_files = [], binary=False):

    if binary:
        # npz
        graph_list = read_binary_graph_raw(raw_dir, add_inverse_edge)
    else:
        # csv
        graph_list = read_csv_graph_raw(raw_dir, add_inverse_edge, additional_node_files = additional_node_files, additional_edge_files = additional_edge_files)
        
    dgl_graph_list = []

    print('Converting graphs into DGL objects...')
    
    for graph in tqdm(graph_list):

        src, dst = torch.from_numpy(graph['edge_index'])
        num_nodes = graph['num_nodes']

        if num_nodes == 1:  # some graphs have one node
            A_ = torch.tensor(1.).view(1, 1)
        else:
            A = torch.zeros([num_nodes, num_nodes], dtype=torch.float)
            A[src, dst] = 1.0
            for i in range(num_nodes):
                A[i, i] = 1.0
            deg = torch.sum(A, axis=0).squeeze() ** -0.5
            D = torch.diag(deg)
            A_ = D @ A @ D
        e, u = torch.linalg.eigh(A_)

        fully_connected = torch.ones([num_nodes, num_nodes], dtype=torch.float).nonzero(as_tuple=True)
        g = dgl.graph(fully_connected, num_nodes = num_nodes)

        g.ndata['e'] = e
        g.ndata['u'] = u

        if graph['node_feat'] is not None:
            g.ndata['feat'] = torch.from_numpy(graph['node_feat'])

        if graph['edge_feat'] is not None:
            edge_idx = torch.stack([src, dst], dim=0)
            edge_attr = torch.from_numpy(graph['edge_feat']) + 1  # for padding

            if len(edge_attr.shape) == 1:
                edge_attr_dense = to_dense_adj(edge_idx, edge_attr=edge_attr.unsqueeze(-1)).squeeze(0).squeeze(-1).view(-1)
            else:
                if edge_attr.size(0) == 0:
                    edge_attr_dense = torch.zeros([num_nodes ** 2, edge_attr.size(1)]).long()  # for graphs without edge
                else:
                    edge_attr_dense = to_dense_adj(edge_idx, edge_attr=edge_attr, max_num_nodes=num_nodes).squeeze(0).view(-1, edge_attr.shape[-1])

            g.edata['feat'] = edge_attr_dense

        dgl_graph_list.append(g)

    return dgl_graph_list
//  [python]
def test_sqrt_inv_matmul_no_lhs(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        with linear_operator.settings.max_cg_iterations(200):
            sqrt_inv_matmul_res = linear_operator.sqrt_inv_matmul(linear_op, rhs)
        evals, evecs = torch.symeig(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = matrix_inv_root @ rhs_copy

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])

        # Perform backward pass
        sqrt_inv_matmul_grad = torch.randn_like(sqrt_inv_matmul_res)
        ((sqrt_inv_matmul_res * sqrt_inv_matmul_grad).sum()).backward()
        ((sqrt_inv_matmul_actual * sqrt_inv_matmul_grad).sum()).backward()

        # Check grads
        self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
            if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def test_sqrt_inv_matmul_no_lhs(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        with linear_operator.settings.max_cg_iterations(200):
            sqrt_inv_matmul_res = linear_operator.sqrt_inv_matmul(linear_op, rhs)
        evals, evecs = torch.linalg.eigh(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = matrix_inv_root @ rhs_copy

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])

        # Perform backward pass
        sqrt_inv_matmul_grad = torch.randn_like(sqrt_inv_matmul_res)
        ((sqrt_inv_matmul_res * sqrt_inv_matmul_grad).sum()).backward()
        ((sqrt_inv_matmul_actual * sqrt_inv_matmul_grad).sum()).backward()

        # Check grads
        self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
            if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def pca_transform(X: torch.Tensor, k: int = 'infer', center: bool = True,
        scale: bool = True):
    """Transforms data to k-dimensional space with identity covariance.

    Use this if you have a wide (p > N) matrix and you only need whitened
    points in the reduces space (an N by k matrix). For the full
    inference procedure use mccfa.pca().

    Args:
        X: Two-dimesional (N x p) torch.tensor.
        k: Integer or 'infer'. Number of pcs to keep. Default is to
          infer using the marchenko pasteur cutoff.
        center: bool. True to mean-center columns of X.
        scale: bool. True to variance-scale the columns of X to 1.0.
    Returns:
        An N by k matrix with X.T @ X / N = I_k.
    """
    N, D = X.shape
    if N > D: raise ValueError('Matrix is tall.')

    if center | scale:
        X = torch.from_numpy(preprocessing.scale(
            X, with_mean=center, with_std=scale))
    mp_lower_bound = 1 + np.sqrt(D / N)

    gram_mat = (X @ X.T) / N
    vals, vecs = torch.eig(gram_mat, eigenvectors=True)
    vals = torch.diag(vals[:,0])
    S = torch.sqrt(torch.abs(torch.flip(vals, dims = [0])))
    A = torch.flip(vecs, dims=[1])
    mp_dim = sum(S > mp_lower_bound)

    if k == 'infer': k = mp_dim
    U = A[:, 0:k]
    return np.sqrt(N) * U
//  [python]
def pca_transform(X: torch.Tensor, k: int = 'infer', center: bool = True,
        scale: bool = True):
    """Transforms data to k-dimensional space with identity covariance.

    Use this if you have a wide (p > N) matrix and you only need whitened
    points in the reduces space (an N by k matrix). For the full
    inference procedure use mccfa.pca().

    Args:
        X: Two-dimesional (N x p) torch.tensor.
        k: Integer or 'infer'. Number of pcs to keep. Default is to
          infer using the marchenko pasteur cutoff.
        center: bool. True to mean-center columns of X.
        scale: bool. True to variance-scale the columns of X to 1.0.
    Returns:
        An N by k matrix with X.T @ X / N = I_k.
    """
    N, D = X.shape
    if N > D: raise ValueError('Matrix is tall.')

    if center | scale:
        X = torch.from_numpy(preprocessing.scale(
            X, with_mean=center, with_std=scale))
    mp_lower_bound = 1 + np.sqrt(D / N)

    gram_mat = (X @ X.T) / N
    vals, vecs = torch.linalg.eigh(gram_mat)
    S = torch.sqrt(torch.abs(torch.flip(vals, dims = [0])))
    A = torch.flip(vecs, dims=[1])
    mp_dim = sum(S > mp_lower_bound)

    if k == 'infer': k = mp_dim
    U = A[:, 0:k]
    return np.sqrt(N) * U
//  [python]
def test_eigh(self):
        dtypes = {"double": torch.double, "float": torch.float}
        for name, dtype in dtypes.items():
            tolerances = self.tolerances["symeig"][name]

            linear_op = self.create_linear_op().detach().requires_grad_(True)
            linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
            evaluated = self.evaluate_linear_op(linear_op_copy)

            # Perform forward pass
            with linalg_dtypes(dtype):
                evals_unsorted, evecs_unsorted = torch.symeig(linear_op)
                evecs_unsorted = evecs_unsorted.to_dense()

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

            evals_actual, evecs_actual = torch.symeig(evaluated.type(dtype))
            evals_actual = evals_actual.to(dtype=evaluated.dtype)
            evecs_actual = evecs_actual.to(dtype=evaluated.dtype)

            # Check forward pass
            self.assertAllClose(evals, evals_actual, **tolerances)
            lt_from_eigendecomp = evecs @ torch.diag_embed(evals) @ evecs.mT
            self.assertAllClose(lt_from_eigendecomp, evaluated, **tolerances)

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

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

            # Check grads if there were no repeated evals
            if not any_evals_repeated:
                for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, **tolerances)
//  [python]
def test_eigh(self):
        dtypes = {"double": torch.double, "float": torch.float}
        for name, dtype in dtypes.items():
            tolerances = self.tolerances["symeig"][name]

            linear_op = self.create_linear_op().detach().requires_grad_(True)
            linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
            evaluated = self.evaluate_linear_op(linear_op_copy)

            # Perform forward pass
            with linalg_dtypes(dtype):
                evals_unsorted, evecs_unsorted = torch.linalg.eigh(linear_op)
                evecs_unsorted = evecs_unsorted.to_dense()

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

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

            # Check forward pass
            self.assertAllClose(evals, evals_actual, **tolerances)
            lt_from_eigendecomp = evecs @ torch.diag_embed(evals) @ evecs.mT
            self.assertAllClose(lt_from_eigendecomp, evaluated, **tolerances)

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

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

            # Check grads if there were no repeated evals
            if not any_evals_repeated:
                for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, **tolerances)
//  [python]
def test_sqrt_inv_matmul(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        lhs = torch.randn(*linear_op.shape[:-2], 2, linear_op.size(-1)).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        lhs_copy = lhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        with linear_operator.settings.max_cg_iterations(200):
            sqrt_inv_matmul_res, inv_quad_res = linear_operator.sqrt_inv_matmul(linear_op, rhs, lhs)
        evals, evecs = torch.symeig(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = lhs_copy @ matrix_inv_root @ rhs_copy
        inv_quad_actual = (lhs_copy @ matrix_inv_root).pow(2).sum(dim=-1)

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(inv_quad_res, inv_quad_actual, **self.tolerances["sqrt_inv_matmul"])

        # Perform backward pass
        sqrt_inv_matmul_grad = torch.randn_like(sqrt_inv_matmul_res)
        inv_quad_grad = torch.randn_like(inv_quad_res)
        ((sqrt_inv_matmul_res * sqrt_inv_matmul_grad).sum() + (inv_quad_res * inv_quad_grad).sum()).backward()
        ((sqrt_inv_matmul_actual * sqrt_inv_matmul_grad).sum() + (inv_quad_actual * inv_quad_grad).sum()).backward()

        # Check grads
        self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(lhs.grad, lhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
            if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def test_sqrt_inv_matmul(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        lhs = torch.randn(*linear_op.shape[:-2], 2, linear_op.size(-1)).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        lhs_copy = lhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        with linear_operator.settings.max_cg_iterations(200):
            sqrt_inv_matmul_res, inv_quad_res = linear_operator.sqrt_inv_matmul(linear_op, rhs, lhs)
        evals, evecs = torch.linalg.eigh(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = lhs_copy @ matrix_inv_root @ rhs_copy
        inv_quad_actual = (lhs_copy @ matrix_inv_root).pow(2).sum(dim=-1)

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(inv_quad_res, inv_quad_actual, **self.tolerances["sqrt_inv_matmul"])

        # Perform backward pass
        sqrt_inv_matmul_grad = torch.randn_like(sqrt_inv_matmul_res)
        inv_quad_grad = torch.randn_like(inv_quad_res)
        ((sqrt_inv_matmul_res * sqrt_inv_matmul_grad).sum() + (inv_quad_res * inv_quad_grad).sum()).backward()
        ((sqrt_inv_matmul_actual * sqrt_inv_matmul_grad).sum() + (inv_quad_actual * inv_quad_grad).sum()).backward()

        # Check grads
        self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(lhs.grad, lhs_copy.grad, **self.tolerances["sqrt_inv_matmul"])
        for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
            if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def _ppca(X, d):
    """Produces a facorization X = W @ W.T + s2*I for symmetric X.

    Args:
      X: symmetric p x p matrix.
      d: number of dimensions to keep.
    """
    vals, vecs = torch.symeig(X)
    vals = torch.abs(torch.flip(vals, dims = [0]))
    vecs = torch.flip(vecs, dims=[1])
    W = vecs[:, 0:d] * torch.sqrt(vals[0:d])
    s2 = torch.mean(vals[d:])
    return W, s2, vals
//  [python]
def _ppca(X, d):
    """Produces a facorization X = W @ W.T + s2*I for symmetric X.

    Args:
      X: symmetric p x p matrix.
      d: number of dimensions to keep.
    """
    vals, vecs = torch.linalg.eigh(X)
    vals = torch.abs(torch.flip(vals, dims = [0]))
    vecs = torch.flip(vecs, dims=[1])
    W = vecs[:, 0:d] * torch.sqrt(vals[0:d])
    s2 = torch.mean(vals[d:])
    return W, s2, vals
//  [python]
def compute_diag_trace(cov: torch.Tensor) -> float:
    """
    Computes the trace of the covariance matrix diagonal matrix
    :param cov: the covariannce matrix as torch tensor
    :return: the trace
    """
    eig_vals, eigen_space = torch.symeig(cov, eigenvectors=True)
    eig_vals[eig_vals < 0] = 0
    return eig_vals.sum().unsqueeze(dim=0).cpu().numpy()[0]
//  [python]
def compute_diag_trace(cov: torch.Tensor) -> float:
    """
    Computes the trace of the covariance matrix diagonal matrix
    :param cov: the covariannce matrix as torch tensor
    :return: the trace
    """
    eig_vals, eigen_space = torch.linalg.eigh(cov, UPLO='U')
    eig_vals[eig_vals < 0] = 0
    return eig_vals.sum().unsqueeze(dim=0).cpu().numpy()[0]
//  [python]
def compute_cov_mat(patches,rank):

    with torch.no_grad():

        # -- cov mat --
        bsize,num,pdim = patches.shape
        covMat = torch.matmul(patches.transpose(2,1),patches)
        covMat /= num

        # -- eigen stuff --
        eigVals,eigVecs = torch.symeig(covMat,eigenvectors=True)
        eigVals= torch.flip(eigVals,dims=(1,))
        eigVecs = torch.flip(eigVecs,dims=(2,))[...,:rank]

    return covMat,eigVals,eigVecs
//  [python]
def compute_cov_mat(patches,rank):

    with torch.no_grad():

        # -- cov mat --
        bsize,num,pdim = patches.shape
        covMat = torch.matmul(patches.transpose(2,1),patches)
        covMat /= num

        # -- eigen stuff --
        eigVals,eigVecs = torch.linalg.eigh(covMat)
        eigVals= torch.flip(eigVals,dims=(1,))
        eigVecs = torch.flip(eigVecs,dims=(2,))[...,:rank]

    return covMat,eigVals,eigVecs
//  [python]
def compute_intrinsic_dimensionality(cov: torch.Tensor,
                                     thresh: float = .99) -> int:
    """
    Compute the intrinsic dimensionality based on the covariance matrix
    :param cov: the covariance matrix as a torch tensor
    :param thresh: delta value; the explained variance of the covariance matrix
    :return: The intrinsic dimensionality; an integer value greater than zero
    """
    eig_vals, eigen_space = torch.symeig(cov, eigenvectors=True)
    eig_vals, idx = eig_vals.sort(descending=True)
    eig_vals[eig_vals < 0] = 0
    percentages = eig_vals.cumsum(0) / eig_vals.sum()
    eigen_space = eigen_space[:, percentages < thresh]
    if eigen_space.shape[1] == 0:
        eigen_space = eigen_space[:, :1]
    elif thresh - (percentages[percentages < thresh][-1]) > 0.02:
        eigen_space = eigen_space[:, :eigen_space.shape[1] + 1]
    return eigen_space.shape[1]
//  [python]
def compute_intrinsic_dimensionality(cov: torch.Tensor,
                                     thresh: float = .99) -> int:
    """
    Compute the intrinsic dimensionality based on the covariance matrix
    :param cov: the covariance matrix as a torch tensor
    :param thresh: delta value; the explained variance of the covariance matrix
    :return: The intrinsic dimensionality; an integer value greater than zero
    """
    eig_vals, eigen_space = torch.linalg.eigh(cov, UPLO='U')
    eig_vals, idx = eig_vals.sort(descending=True)
    eig_vals[eig_vals < 0] = 0
    percentages = eig_vals.cumsum(0) / eig_vals.sum()
    eigen_space = eigen_space[:, percentages < thresh]
    if eigen_space.shape[1] == 0:
        eigen_space = eigen_space[:, :1]
    elif thresh - (percentages[percentages < thresh][-1]) > 0.02:
        eigen_space = eigen_space[:, :eigen_space.shape[1] + 1]
    return eigen_space.shape[1]
//  [python]
def pca(X: pd.DataFrame, k: int = 'infer', center: bool = True,
        scale: bool = True, calc_V = True) -> PCARes:
    """Basic PCA implementation.

    This is a basic PCA implementation which is particularly efficient for
    top-k PCA by doing the eigendecomposition of either X X.T / N or
    X.T X / N.

    Args:
        X: n (samples) by d (features) pandas DataFrame.
        k: Integer, 'infer' or 'all'. Number of pcs to keep. Default is to
          infer using the marchenko pasteur cutoff.
        center: bool. True to mean-center columns of X.
        scale: bool. True to variance-scale the columns of X to 1.0.
        calc_V: True to track the PC loadings (right singular vectors)
          of X. Setting to False can save substantial memory if X is very
          wide.
    Returns:
        A PCARes instance.
    """
    sample_names = X.index

    if center | scale:
        X = torch.from_numpy(preprocessing.scale(
            X, with_mean=center, with_std=scale))
    else:
        X = torch.from_numpy(X.values)
    N, D = X.shape
    mp_lower_bound = 1 + np.sqrt(D / N)

    gram_mat = (X @ X.T) / N if D > N else X.T @ X / N
    vals, vecs = torch.eig(gram_mat, eigenvectors=True)
    vals = torch.real(vals[:,0])
    vecs = torch.real(vecs)
    S = torch.sqrt(torch.abs(torch.flip(vals, dims = [0])))
    A = torch.flip(vecs, dims=[1])
    mp_dim = sum(S > mp_lower_bound)

    # TODO(brielin): Throw error if infer and not center/scale.
    if k in ('infer', 'all'):
        k = mp_dim if k == 'infer' else min(N, D)
    S_k = S[0:k]

    if D > N:
        U = A[:, 0:k]
        V = torch.eye(k, dtype=torch.double)
        if calc_V:
            V = X.T @ U / (S_k * np.sqrt(N))
    else:
        V = A[:, 0:k]
        U = X @ V / (S_k * np.sqrt(N))
        V = V if calc_V else torch.eye(k, dtype=torch.double)

    pc_names = ['PC' + str(i+1) for i in range(k)]
    pcs = pd.DataFrame((U * S_k * np.sqrt(N)).numpy(), index=sample_names,
                       columns=pc_names)
    var_exp = pd.Series(S_k**2 / torch.sum(S**2), index=pc_names)
    return PCARes(pcs, var_exp, U, S_k, V, k, N, D, mp_dim)
//  [python]
def pca(X: pd.DataFrame, k: int = 'infer', center: bool = True,
        scale: bool = True, calc_V = True) -> PCARes:
    """Basic PCA implementation.

    This is a basic PCA implementation which is particularly efficient for
    top-k PCA by doing the eigendecomposition of either X X.T / N or
    X.T X / N.

    Args:
        X: n (samples) by d (features) pandas DataFrame.
        k: Integer, 'infer' or 'all'. Number of pcs to keep. Default is to
          infer using the marchenko pasteur cutoff.
        center: bool. True to mean-center columns of X.
        scale: bool. True to variance-scale the columns of X to 1.0.
        calc_V: True to track the PC loadings (right singular vectors)
          of X. Setting to False can save substantial memory if X is very
          wide.
    Returns:
        A PCARes instance.
    """
    sample_names = X.index

    if center | scale:
        X = torch.from_numpy(preprocessing.scale(
            X, with_mean=center, with_std=scale))
    else:
        X = torch.from_numpy(X.values)
    N, D = X.shape
    mp_lower_bound = 1 + np.sqrt(D / N)

    gram_mat = (X @ X.T) / N if D > N else X.T @ X / N
    vals, vecs = torch.linalg.eigh(gram_mat)
    S = torch.sqrt(torch.abs(torch.flip(vals, dims = [0])))
    A = torch.flip(vecs, dims=[1])
    mp_dim = sum(S > mp_lower_bound)

    # TODO(brielin): Throw error if infer and not center/scale.
    if k in ('infer', 'all'):
        k = mp_dim if k == 'infer' else min(N, D)
    S_k = S[0:k]

    if D > N:
        U = A[:, 0:k]
        V = torch.eye(k, dtype=torch.double)
        if calc_V:
            V = X.T @ U / (S_k * np.sqrt(N))
    else:
        V = A[:, 0:k]
        U = X @ V / (S_k * np.sqrt(N))
        V = V if calc_V else torch.eye(k, dtype=torch.double)

    pc_names = ['PC' + str(i+1) for i in range(k)]
    pcs = pd.DataFrame((U * S_k * np.sqrt(N)).numpy(), index=sample_names,
                       columns=pc_names)
    var_exp = pd.Series(S_k**2 / torch.sum(S**2), index=pc_names)
    return PCARes(pcs, var_exp, U, S_k, V, k, N, D, mp_dim)
//  [python]
def auto_orient_and_center_poses(
    i_train,
    poses: TensorType["num_poses":..., 4, 4],
    method: Literal["pca", "up", "none"] = "up",
    center_poses: bool = True,
) -> TensorType["num_poses":..., 3, 4]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        The oriented poses.
    """

    translation = poses[i_train][..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        w, v = torch.symeig(translation_diff.T @ translation_diff, eigenvectors=True)
        eigvec = torch.flip(v, dims=(-1,))

        if torch.linalg.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[i_train][:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def auto_orient_and_center_poses(
    i_train,
    poses: TensorType["num_poses":..., 4, 4],
    method: Literal["pca", "up", "none"] = "up",
    center_poses: bool = True,
) -> TensorType["num_poses":..., 3, 4]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        The oriented poses.
    """

    translation = poses[i_train][..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.linalg.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[i_train][:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def eig_decomposition(matrix, epsilon=1e-6, mode='regularize'):
    """ This method can be applied to do the eig-decomposition for a rank deficient hermetian matrix,
    this method will be further used to estimate koopman matrix.

    Parameters
    ----------
    matrix : torch.Tensor
        The hermitian matrix: specifically, the covariance matrix.
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    (eigval, eigvec) : Tuple[torch.Tensor, torch.Tensor]
        Eigenvalues and eigenvectors.
    """

    if mode == 'regularize':
        identity = torch.eye(matrix.shape[0], dtype=matrix.dtype, device=matrix.device)
        matrix = matrix + epsilon * identity
        eigval, eigvec = torch.symeig(matrix)
        eigval = torch.abs(eigval)
        eigvec = eigvec.transpose(0, 1)  # row -> column

    elif mode == 'trunc':
        eigval, eigvec = torch.symeig(matrix)
        eigvec = eigvec.transpose(0, 1)
        mask = eigval > epsilon
        eigval = eigval[mask]
        eigvec = eigvec[mask]

    else:
        raise ValueError('Mode is not included')

    return eigval, eigvec
//  [python]
def eig_decomposition(matrix, epsilon=1e-6, mode='regularize'):
    """ This method can be applied to do the eig-decomposition for a rank deficient hermetian matrix,
    this method will be further used to estimate koopman matrix.

    Parameters
    ----------
    matrix : torch.Tensor
        The hermitian matrix: specifically, the covariance matrix.
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    (eigval, eigvec) : Tuple[torch.Tensor, torch.Tensor]
        Eigenvalues and eigenvectors.
    """

    if mode == 'regularize':
        identity = torch.eye(matrix.shape[0], dtype=matrix.dtype, device=matrix.device)
        matrix = matrix + epsilon * identity
        eigval, eigvec = torch.linalg.eigh(matrix)
        eigval = torch.abs(eigval)
        eigvec = eigvec.transpose(0, 1)  # row -> column

    elif mode == 'trunc':
        eigval, eigvec = torch.linalg.eigh(matrix)
        eigvec = eigvec.transpose(0, 1)
        mask = eigval > epsilon
        eigval = eigval[mask]
        eigvec = eigvec[mask]

    else:
        raise ValueError('Mode is not included')

    return eigval, eigvec