jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def eig(self):
        """
        Returns the eigenvalues (optionally eigenvectors) of the linear map used in matrix multiplication.

        :return: (torch.Tensor) Vector of eigenvalues, optionally a tuple including a matrix of eigenvectors.
        """
        return torch.linalg.eig(self.effective_W())
//  [python]
def test_eig_compare_backends(self, device, dtype):
        def run_test(shape, *, symmetric=False):
            from torch.testing._internal.common_utils import random_symmetric_matrix

            if not dtype.is_complex and symmetric:
                # for symmetric real-valued inputs eigenvalues and eigenvectors have imaginary part equal to zero
                a = random_symmetric_matrix(shape[-1], *shape[:-2], dtype=dtype, device=device)
            else:
                a = make_tensor(shape, dtype=dtype, device=device)

            actual = torch.eig(a)

            complementary_device = 'cpu'

            # compare with CPU
            expected = torch.eig(a.to(complementary_device))
            self.assertEqual(expected[0], actual[0])
            self.assertEqual(expected[1], actual[1])

        shapes = [(0, 0),  # Empty matrix
                  (5, 5),  # Single matrix
                  (0, 0, 0), (0, 5, 5),  # Zero batch dimension tensors
                  (2, 5, 5),  # 3-dim tensors
                  (2, 1, 5, 5)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
            run_test(shape, symmetric=True)
//  [python]
def test_eig_compare_backends(self, device, dtype):
        def run_test(shape, *, symmetric=False):
            from torch.testing._internal.common_utils import random_symmetric_matrix

            if not dtype.is_complex and symmetric:
                # for symmetric real-valued inputs eigenvalues and eigenvectors have imaginary part equal to zero
                a = random_symmetric_matrix(shape[-1], *shape[:-2], dtype=dtype, device=device)
            else:
                a = make_tensor(shape, dtype=dtype, device=device)

            actual = torch.linalg.eig(a)

            complementary_device = 'cpu'

            # compare with CPU
            expected = torch.linalg.eig(a.to(complementary_device))
            self.assertEqual(expected[0], actual[0])
            self.assertEqual(expected[1], actual[1])

        shapes = [(0, 0),  # Empty matrix
                  (5, 5),  # Single matrix
                  (0, 0, 0), (0, 5, 5),  # Zero batch dimension tensors
                  (2, 5, 5),  # 3-dim tensors
                  (2, 1, 5, 5)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
            run_test(shape, symmetric=True)
//  [python]
def eigenvalues(self):
        """ Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k. """
        return torch.eig(self.BA)[0][:, 0]
//  [python]
def eigenvalues(self):
        """ Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k. """
        return torch.linalg.eig(self.BA).eigenvalues
//  [python]
def test_eig_numpy(self, device, dtype):
        def run_test(shape, *, symmetric=False):
            from torch.testing._internal.common_utils import random_symmetric_matrix

            if not dtype.is_complex and symmetric:
                # for symmetric real-valued inputs eigenvalues and eigenvectors have imaginary part equal to zero
                # unlike NumPy the result is not cast to float32 or float64 dtype in this case
                a = random_symmetric_matrix(shape[-1], *shape[:-2], dtype=dtype, device=device)
            else:
                a = make_tensor(shape, dtype=dtype, device=device)

            actual = torch.eig(a)

            # compare with NumPy
            # the eigenvalues are not necessarily ordered
            # so order of NumPy and PyTorch can be different
            expected = np.linalg.eig(a.cpu().numpy())

            # sort NumPy output
            ind = np.argsort(expected[0], axis=-1)[::-1]
            expected = (np.take_along_axis(expected[0], ind, axis=-1), np.take_along_axis(expected[1], ind[:, None], axis=-1))

            # sort PyTorch output
            # torch.argsort doesn't work with complex inputs, NumPy sorting on CPU is used instead
            # RuntimeError: _th_sort not supported on CUDAType for ComplexDouble
            # RuntimeError: "sorting_kernel_method_name" not implemented for 'ComplexDouble'
            ind = np.argsort(actual[0].cpu().numpy(), axis=-1)[::-1]
            actual_np = [x.cpu().numpy() for x in actual]
            sorted_actual = (
                np.take_along_axis(actual_np[0], ind, axis=-1),
                np.take_along_axis(actual_np[1], ind[:, None], axis=-1))

            self.assertEqual(expected[0], sorted_actual[0])
            self.assertEqual(abs(expected[1]), abs(sorted_actual[1]))

        shapes = [(0, 0),  # Empty matrix
                  (5, 5),  # Single matrix
                  (0, 0, 0), (0, 5, 5),  # Zero batch dimension tensors
                  (2, 5, 5),  # 3-dim tensors
                  (2, 1, 5, 5)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
            run_test(shape, symmetric=True)
//  [python]
def test_eig_numpy(self, device, dtype):
        def run_test(shape, *, symmetric=False):
            from torch.testing._internal.common_utils import random_symmetric_matrix

            if not dtype.is_complex and symmetric:
                # for symmetric real-valued inputs eigenvalues and eigenvectors have imaginary part equal to zero
                # unlike NumPy the result is not cast to float32 or float64 dtype in this case
                a = random_symmetric_matrix(shape[-1], *shape[:-2], dtype=dtype, device=device)
            else:
                a = make_tensor(shape, dtype=dtype, device=device)

            actual = torch.linalg.eig(a)

            # compare with NumPy
            # the eigenvalues are not necessarily ordered
            # so order of NumPy and PyTorch can be different
            expected = np.linalg.eig(a.cpu().numpy())

            # sort NumPy output
            ind = np.argsort(expected[0], axis=-1)[::-1]
            expected = (np.take_along_axis(expected[0], ind, axis=-1), np.take_along_axis(expected[1], ind[:, None], axis=-1))

            # sort PyTorch output
            # torch.argsort doesn't work with complex inputs, NumPy sorting on CPU is used instead
            # RuntimeError: _th_sort not supported on CUDAType for ComplexDouble
            # RuntimeError: "sorting_kernel_method_name" not implemented for 'ComplexDouble'
            ind = np.argsort(actual[0].cpu().numpy(), axis=-1)[::-1]
            actual_np = [x.cpu().numpy() for x in actual]
            sorted_actual = (
                np.take_along_axis(actual_np[0], ind, axis=-1),
                np.take_along_axis(actual_np[1], ind[:, None], axis=-1))

            self.assertEqual(expected[0], sorted_actual[0])
            self.assertEqual(abs(expected[1]), abs(sorted_actual[1]))

        shapes = [(0, 0),  # Empty matrix
                  (5, 5),  # Single matrix
                  (0, 0, 0), (0, 5, 5),  # Zero batch dimension tensors
                  (2, 5, 5),  # 3-dim tensors
                  (2, 1, 5, 5)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
            run_test(shape, symmetric=True)
//  [python]
def make_positive_definite(X, eps=1.0e-06):
    assert eps >= 0.
    Lmd, V = torch.eig(X)
    Lmd.real = torch.clamp(Lmd.real, min=eps)
    return V.mm(torch.diag(Lmd)).mm(V.transpose(0, 1)).real
//  [python]
def make_positive_definite(X, eps=1.0e-06):
    assert eps >= 0.
    Lmd, V = torch.linalg.eig(X)
    Lmd.real = torch.clamp(Lmd.real, min=eps)
    return V.mm(torch.diag(Lmd)).mm(V.transpose(0, 1)).real
//  [python]
def get_log_like_gradient_and_hessian(self, X, eps=1e-2, compute_hessian=False):
        log_pi = self.V(X)
        if compute_hessian:
            self.hessian_counter += 1
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.autograd.functional.hessian(self.V, X, create_graph=True)
            lamda, Q = torch.linalg.eig(H)
            lamda, Q = torch.real(lamda), torch.real(Q)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / lamda_prime
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def eig(
    x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    if not torch.is_complex(x):
        x = x.to(torch.complex128)
    return torch.eig(x)
//  [python]
def eig(
    x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    if not torch.is_complex(x):
        x = x.to(torch.complex128)
    return torch.linalg.eig(x)
//  [python]
def _setup_step(self, mode='dense'):
        """Set up dA, dB, dC discretized parameters for stepping."""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C)) # (H C N)
        if self.l_kernel.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.l_kernel.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("DPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def _setup_step(self, mode='dense'):
        """Set up dA, dB, dC discretized parameters for stepping."""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C)) # (H C N)
        if self.l_kernel.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.l_kernel.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("DPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def color_histogram_matching(
    content: torch.Tensor,
    style: torch.Tensor,
) -> torch.Tensor:
    n, c, _, _ = style.shape

    x, y = torch.flatten(style, start_dim=2), torch.flatten(content, start_dim=2)

    mu_x, mu_y = x.mean(2, keepdim=True), y.mean(2, keepdim=True)
    cov_x = torch.bmm(x - mu_x, (x - mu_x).transpose(1, 2)) / x.numel()
    cov_y = torch.bmm(y - mu_y, (y - mu_y).transpose(1, 2)) / y.numel()

    eta_x, eig_x = torch.symeig(cov_x, eigenvectors=True)
    eta_y, eig_y = torch.symeig(cov_y, eigenvectors=True)

    cov_half_x = torch.bmm(
        torch.bmm(eig_x, eta_x.diag_embed().sqrt()),
        eig_x.transpose(1, 2),
    )
    cov_half_y = torch.bmm(
        torch.bmm(eig_y, eta_y.diag_embed().sqrt()),
        eig_y.transpose(1, 2),
    )

    A = torch.bmm(cov_half_y, cov_half_x.inverse()).real

    print(A)

    return (
        torch.bmm(A, x).reshape(style.shape)
        + mu_y.unsqueeze(3)
        - torch.bmm(A, mu_x).unsqueeze(3)
    )
//  [python]
def color_histogram_matching(
    content: torch.Tensor,
    style: torch.Tensor,
) -> torch.Tensor:
    n, c, _, _ = style.shape

    x, y = torch.flatten(style, start_dim=2), torch.flatten(content, start_dim=2)

    mu_x, mu_y = x.mean(2, keepdim=True), y.mean(2, keepdim=True)
    cov_x = torch.bmm(x - mu_x, (x - mu_x).transpose(1, 2)) / x.numel()
    cov_y = torch.bmm(y - mu_y, (y - mu_y).transpose(1, 2)) / y.numel()

    eta_x, eig_x = torch.linalg.eig(cov_x)
    eta_y, eig_y = torch.linalg.eig(cov_y)

    cov_half_x = torch.bmm(
        torch.bmm(eig_x, eta_x.diag_embed().sqrt()),
        eig_x.transpose(1, 2),
    )
    cov_half_y = torch.bmm(
        torch.bmm(eig_y, eta_y.diag_embed().sqrt()),
        eig_y.transpose(1, 2),
    )

    A = torch.bmm(cov_half_y, cov_half_x.inverse()).real

    print(A)

    return (
        torch.bmm(A, x).reshape(style.shape)
        + mu_y.unsqueeze(3)
        - torch.bmm(A, mu_x).unsqueeze(3)
    )
//  [python]
def __calc_score(jacobian: Tensor):
        """Calculate the score for a Jacobian."""
        correlations = ObjectiveNASWOT.__corrcoef(jacobian)
        eigen_values, _ = torch.symeig(correlations)
        # Only consider the real valued part, imaginary part is rounding
        # artefact
        eigen_values = eigen_values.real
        # Needed for numerical stability. 1e-4 instead of 1e-5 in reference
        # as the torchs eigenvalue solver on GPU
        # seems to have bigger rounding errors than numpy, resulting in
        # slightly larger negative Eigenvalues
        k = 1e-4
        v = -torch.sum(torch.log(eigen_values + k) + 1.0 / (eigen_values + k))
        return v
//  [python]
def __calc_score(jacobian: Tensor):
        """Calculate the score for a Jacobian."""
        correlations = ObjectiveNASWOT.__corrcoef(jacobian)
        eigen_values, _ = torch.linalg.eig(correlations)
        # Only consider the real valued part, imaginary part is rounding
        # artefact
        eigen_values = eigen_values.real
        # Needed for numerical stability. 1e-4 instead of 1e-5 in reference
        # as the torchs eigenvalue solver on GPU
        # seems to have bigger rounding errors than numpy, resulting in
        # slightly larger negative Eigenvalues
        k = 1e-4
        v = -torch.sum(torch.log(eigen_values + k) + 1.0 / (eigen_values + k))
        return v
//  [python]
def forward(self, x, y):
        def kernel_matrix(data, kernel="linear", sigma=None):
            if kernel == "linear":
                return torch.matmul(data, data.t())
            elif kernel == "rbf":
                square = torch.square(data)
                sum_square = torch.sum(square, dim=-1, keepdim=True)
                distances = (
                    sum_square + sum_square.t() - 2 * torch.matmul(data, data.t())
                )
                return torch.exp(-distances / (2 * sigma**2))
            else:
                raise ValueError("Invalid kernel type.")

        Kx = kernel_matrix(x, kernel=self.kernel, sigma=self.sigma)
        Ky = kernel_matrix(y, kernel=self.kernel, sigma=self.sigma)

        n = x.size(0)
        I = torch.eye(n).to(x.device)

        L, R = torch.inverse(Kx + self.reg * I), torch.inverse(Ky + self.reg * I)
        L_inv, R_inv = torch.inverse(L), torch.inverse(R)

        M = torch.matmul(L_inv, torch.matmul(Kx, torch.matmul(Ky, R_inv)))

        eigenvalues = torch.symeig(M, eigenvectors=False).eigenvalues
        kcca_corr = torch.sum(torch.sqrt(torch.clamp(eigenvalues.real, min=0)))

        return kcca_corr
//  [python]
def forward(self, x, y):
        def kernel_matrix(data, kernel="linear", sigma=None):
            if kernel == "linear":
                return torch.matmul(data, data.t())
            elif kernel == "rbf":
                square = torch.square(data)
                sum_square = torch.sum(square, dim=-1, keepdim=True)
                distances = (
                    sum_square + sum_square.t() - 2 * torch.matmul(data, data.t())
                )
                return torch.exp(-distances / (2 * sigma**2))
            else:
                raise ValueError("Invalid kernel type.")

        Kx = kernel_matrix(x, kernel=self.kernel, sigma=self.sigma)
        Ky = kernel_matrix(y, kernel=self.kernel, sigma=self.sigma)

        n = x.size(0)
        I = torch.eye(n).to(x.device)

        L, R = torch.inverse(Kx + self.reg * I), torch.inverse(Ky + self.reg * I)
        L_inv, R_inv = torch.inverse(L), torch.inverse(R)

        M = torch.matmul(L_inv, torch.matmul(Kx, torch.matmul(Ky, R_inv)))

        eigenvalues, _ = torch.linalg.eig(M)
        kcca_corr = torch.sum(torch.sqrt(torch.clamp(eigenvalues.real, min=0)))

        return kcca_corr
//  [python]
def forward(ctx,x):
        ctx.input = x
        eigval, eigvec = torch.eig(x)
        ctx.eigval = eigval.cpu()
        ctx.eigvec = eigvec.cpu()
        return eigval, eigvec
//  [python]
def forward(ctx,x):
        ctx.input = x
        eigval, eigvec = torch.linalg.eig(x)
        ctx.eigval = eigval.cpu()
        ctx.eigvec = eigvec.cpu()
        return eigval, eigvec
//  [python]
def nplr(measure, n, rank=1, dtype=torch.float):
    """
    Return w, p, q, v, b such that
    (w - p q^*, b) is unitarily equivalent to the original HiPPO a, b by the matrix v
    i.e. a = v[w - p q^*]v^*, b = v b
    Args:
        measure (str): the type of measure
            legt - Legendre (translated)
            legs - Legendre (scaled)
            glagt - generalized Laguerre (translated)
            lagt, tlagt - previous versions of (tilted) Laguerre with slightly different normalization
        n (int): dimension
        rank (int): rank
        dtype (str): data type
    """
    a, b = transition(measure, n)
    a = torch.as_tensor(a, dtype=dtype)  # (n, n)
    b = torch.as_tensor(b, dtype=dtype)[:, 0]  # (n,)

    p = rank_correction(measure, n, rank=rank, dtype=dtype)
    ap = a + torch.sum(p.unsqueeze(-2)*p.unsqueeze(-1), dim=-3)
    w, v = torch.eig(ap)  # (..., n) (..., n, n)
    # v w v^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    v = v[..., 0::2].contiguous()

    v_inv = v.conj().transpose(-1, -2)

    b = contract('ij, j -> i', v_inv, b.to(v))  # v^* b
    p = contract('ij, ...j -> ...i', v_inv, p.to(v))  # v^* p

    return w, p, p, b, v
//  [python]
def nplr(measure, n, rank=1, dtype=torch.float):
    """
    Return w, p, q, v, b such that
    (w - p q^*, b) is unitarily equivalent to the original HiPPO a, b by the matrix v
    i.e. a = v[w - p q^*]v^*, b = v b
    Args:
        measure (str): the type of measure
            legt - Legendre (translated)
            legs - Legendre (scaled)
            glagt - generalized Laguerre (translated)
            lagt, tlagt - previous versions of (tilted) Laguerre with slightly different normalization
        n (int): dimension
        rank (int): rank
        dtype (str): data type
    """
    a, b = transition(measure, n)
    a = torch.as_tensor(a, dtype=dtype)  # (n, n)
    b = torch.as_tensor(b, dtype=dtype)[:, 0]  # (n,)

    p = rank_correction(measure, n, rank=rank, dtype=dtype)
    ap = a + torch.sum(p.unsqueeze(-2)*p.unsqueeze(-1), dim=-3)
    w, v = torch.linalg.eig(ap)  # (..., n) (..., n, n)
    # v w v^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    v = v[..., 0::2].contiguous()

    v_inv = v.conj().transpose(-1, -2)

    b = contract('ij, j -> i', v_inv, b.to(v))  # v^* b
    p = contract('ij, ...j -> ...i', v_inv, p.to(v))  # v^* p

    return w, p, p, b, v
//  [python]
def positional_encoding(g, pos_enc_dim):
    """
        Graph positional encoding v/ Laplacian eigenvectors
    """

    # Laplacian
    A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    A = torch.tensor(A.toarray()).float().to("cuda")
    N = torch.diag(g.in_degrees().clip(1) ** -0.5).float().to("cuda")
    L = torch.eye(g.number_of_nodes()).to("cuda") - N.mm(A).mm(N)

    # # Eigenvectors with numpy
    EigVal, EigVec = torch.linalg.eig(L)
    idx = torch.real(EigVal).argsort() # increasing order
    EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
    g.ndata['pos_enc'] = EigVec[:,1:pos_enc_dim+1].cpu().float() 

    # A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    # N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=float)
    # L = sp.eye(g.number_of_nodes()) - N * A * N
    # EigVal, EigVec = sp.linalg.eigs(L, k=pos_enc_dim+1, which='SR', tol=1e-2) # for 40 PEs
    # EigVec = EigVec[:, EigVal.argsort()] # increasing order
    # g.ndata['pos_enc'] = torch.from_numpy(np.real(EigVec[:,1:pos_enc_dim+1])).float()


    return g
//  [python]
def laplace_pe(adj):
    """
        Graph positional encoding v/ Laplacian eigenvectors
    """

    # Laplacian
    edge_attr = adj.long()
    edge_degree = torch.diag_embed(edge_attr.sum(dim=-1))
    edge_laplace=edge_degree-edge_attr

    # 求解特征值与特征向量
    for x in range(len(edge_laplace)):
        if x==0:
            t=torch.eig(edge_laplace[x].float(),eigenvectors=True)[1]
            # t=torch.linalg.eig(edge_laplace[x].float())[1]
        else:
            t=torch.cat([t, torch.eig(edge_laplace[x].float(),eigenvectors=True)[1]],dim=0)
            # t=torch.cat([t, torch.linalg.eig(edge_laplace[x].float())[1]],dim=0)
    feature_matrix=t.reshape(len(adj),t.size(1),t.size(1))           
    return feature_matrix
//  [python]
def laplace_pe(adj):
    """
        Graph positional encoding v/ Laplacian eigenvectors
    """

    # Laplacian
    edge_attr = adj.long()
    edge_degree = torch.diag_embed(edge_attr.sum(dim=-1))
    edge_laplace=edge_degree-edge_attr

    # 求解特征值与特征向量
    for x in range(len(edge_laplace)):
        if x==0:
            # t=torch.eig(edge_laplace[x].float(),eigenvectors=True)[1]
            t=torch.linalg.eig(edge_laplace[x].float())[1]
        else:
            # t=torch.cat([t, torch.eig(edge_laplace[x].float(),eigenvectors=True)[1]],dim=0)
            t=torch.cat([t, torch.linalg.eig(edge_laplace[x].float())[1]],dim=0)
    feature_matrix=t.reshape(len(adj),t.size(1),t.size(1))           
    return feature_matrix
//  [python]
def laplacian_positional_enc(self, adj):
        
        A = adj
        D = torch.diag(torch.count_nonzero(A, dim=-1))
        L = torch.eye(A.shape[0], device=A.device) - D * A * D
        
        EigVal, EigVec = torch.eig(L, eigenvectors=True)
    
        idx = torch.argsort(torch.real(EigVal[:,0]))
        EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
        pos_enc = EigVec[:,1:self.pos_enc_dim + 1]
        
        return pos_enc
//  [python]
def laplacian_positional_enc(self, adj):
        
        A = adj
        D = torch.diag(torch.count_nonzero(A, dim=-1))
        L = torch.eye(A.shape[0], device=A.device) - D * A * D
        
        EigVal, EigVec = torch.linalg.eig(L)
    
        idx = torch.argsort(torch.real(EigVal))
        EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
        pos_enc = EigVec[:,1:self.pos_enc_dim + 1]
        
        return pos_enc
//  [python]
def MKL(A, B):
    """only torch>=1.10

    Args:
        A (torch.tensor): inpute shape:(3,3)
        B (torch.tensor): target shape:(3,3)
    """
    eps = 1e-10
    N = A.shape[0]
    Da2, Ua = torch.eig(A, eigenvectors=True)
    Da2 = Da2[:,0].real.to(A.device)
    Ua = Ua.real.to(A.device)

    Da2[Da2<0] = 0
    Da = (Da2 + eps).sqrt().diag()  # (N,N)

    # U = Ua.T.contiguous()
    C = Da.mm(Ua.T).mm(B).mm(Ua).mm(Da)  # (N,N)

    Dc2_, Uc_ = torch.eig(C.clone(), eigenvectors=True)
    Dc2 = Dc2_[:,0].real.to(A.device)
    Uc = Uc_.real.to(A.device)
    Dc2[Dc2<0] = 0

    Dc = (Dc2 + eps).sqrt().diag()  # (N,N)

    Da_inv = (1/(Da.diag())).diag()  # (N,N)

    T = Ua.mm(Da_inv).mm(Uc).mm(Dc).mm(Uc.T).mm(Da_inv).mm(Ua.T)

    return T
//  [python]
def MKL(A, B):
    """only torch>=1.10

    Args:
        A (torch.tensor): inpute shape:(3,3)
        B (torch.tensor): target shape:(3,3)
    """
    eps = 1e-10
    N = A.shape[0]
    Da2, Ua = torch.linalg.eig(A)
    Da2 = Da2.real.to(A.device)
    Ua = Ua.real.to(A.device)

    Da2[Da2<0] = 0
    Da = (Da2 + eps).sqrt().diag()  # (N,N)

    # U = Ua.T.contiguous()
    C = Da.mm(Ua.T).mm(B).mm(Ua).mm(Da)  # (N,N)

    Dc2_, Uc_ = torch.linalg.eig(C.clone())
    Dc2 = Dc2_.real.to(A.device)
    Uc = Uc_.real.to(A.device)
    Dc2[Dc2<0] = 0

    Dc = (Dc2 + eps).sqrt().diag()  # (N,N)

    Da_inv = (1/(Da.diag())).diag()  # (N,N)

    T = Ua.mm(Da_inv).mm(Uc).mm(Dc).mm(Uc.T).mm(Da_inv).mm(Ua.T)

    return T
//  [python]
def calculate_scaled_laplacian_torch(adj):
    """
    torch version
    """
    D = torch.diag(torch.sum(adj, dim=1))
    L = D - adj

    eigenvalues, _ = torch.eig(L)
    lambda_max = eigenvalues[:,0].real.max()

    return (2 * L) / lambda_max - torch.eye(adj.shape[0]).cuda()
//  [python]
def calculate_scaled_laplacian_torch(adj):
    """
    torch version
    """
    D = torch.diag(torch.sum(adj, dim=1))
    L = D - adj

    eigenvalues, _ = torch.linalg.eig(L)
    lambda_max = eigenvalues.real.max()

    return (2 * L) / lambda_max - torch.eye(adj.shape[0]).cuda()
//  [python]
def get_avg_eigvec_angle_diff(matrix: torch.Tensor) -> float:
    _, eigvecs = torch.eig(matrix)
    avg_angle = 0
    for v1, v2 in itertools.pairwise(eigvecs):
        avg_angle += vec_angle(v1, v2)

    res = avg_angle / len(eigvecs)
    if isinstance(res, complex):
        return res.real
    else:
        return res
//  [python]
def get_avg_eigvec_angle_diff(matrix: torch.Tensor) -> float:
    _, eigvecs = torch.linalg.eig(matrix)
    avg_angle = 0
    for v1, v2 in itertools.pairwise(eigvecs):
        avg_angle += vec_angle(v1, v2)

    res = avg_angle / len(eigvecs)
    if isinstance(res, complex):
        return res.real
    else:
        return res
//  [python]
def modeig_forward_re(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape
    U, S = torch.zeros_like(P, device=P.device), torch.zeros(batch_size,
                                                             channels,
                                                             n,
                                                             dtype=P.dtype,
                                                             device=P.device)
    for i in range(batch_size):
        for j in range(channels):
            if (eig_mode == 'eig'):
                s, U[i, j] = torch.eig(P[i, j], True)
                S[i, j] = s[:, 0]
            elif (eig_mode == 'svd'):
                U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j], 1e-5))
                # U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j].cpu(), 1e-5))
    # U, S = U.cuda(), S.cuda()
    S_fn = op.fn(S, param)
    X = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2, 3))
    return X, U, S, S_fn
//  [python]
def modeig_forward_re(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape
    U, S = torch.zeros_like(P, device=P.device), torch.zeros(batch_size,
                                                             channels,
                                                             n,
                                                             dtype=P.dtype,
                                                             device=P.device)
    for i in range(batch_size):
        for j in range(channels):
            if (eig_mode == 'eig'):
                s, U[i, j] = torch.linalg.eig(P[i, j], True)
                S[i, j] = s[:, 0]
            elif (eig_mode == 'svd'):
                U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j], 1e-5))
                # U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j].cpu(), 1e-5))
    # U, S = U.cuda(), S.cuda()
    S_fn = op.fn(S, param)
    X = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2, 3))
    return X, U, S, S_fn
//  [python]
def eigenvalues(self) -> Float[torch.Tensor, "*leading_dims mdim"]:
        """Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k."""
        return torch.eig(self.BA)[0]
//  [python]
def eigenvalues(self) -> Float[torch.Tensor, "*leading_dims mdim"]:
        """Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k."""
        return torch.linalg.eig(self.BA).eigenvalues
//  [python]
def matrix_logarithm(A):
    L, V = torch.eig(A)
    return (V @ torch.diag_embed(torch.log(L + 1e-20)) @ torch.linalg.inv(V)).real
//  [python]
def matrix_logarithm(A):
    L, V = torch.linalg.eig(A)
    return (V @ torch.diag_embed(torch.log(L + 1e-20)) @ torch.linalg.inv(V)).real
//  [python]
def eig_calc_gpu(L):
    device = torch.device("cuda:0") 
    L=torch.from_numpy(L.toarray()).to(device)

    evals_sn, evects_sn = torch.eig(L, eigenvectors=True)

    evals_sn = evals_sn.detach().cpu().numpy()
    evects_sn = evects_sn.detach().cpu().numpy()
    return evals_sn, evects_sn
//  [python]
def eig_calc_gpu(L):
    device = torch.device("cuda:0") 
    L=torch.from_numpy(L.toarray()).to(device)

    evals_sn, evects_sn = torch.linalg.eig(L)

    evals_sn = evals_sn.detach().cpu().numpy()
    evects_sn = evects_sn.detach().cpu().numpy()
    return evals_sn, evects_sn
//  [python]
def get_weight_infos(models: list[MLP]) -> tuple[torch.Tensor, ...]:
    abs_mean_vals = []
    abs_mean_diffs = []
    eigval_diffs = []
    eigvec_angles = []
    for info in zip(*[model.named_parameters() for model in models]):
        if "weight" not in info[0][0] or info[0][1].ndim != 2:
            continue

        abs_means = torch.tensor([p.abs().mean() for _, p in info])
        abs_mean_vals.append(torch.mean(abs_means))
        abs_mean_diff = (
                (torch.max(abs_means) - torch.min(abs_means))
                / torch.mean(abs_means)  # always positive due to squaring
        )
        abs_mean_diffs.append(abs_mean_diff)

        # Eigval & eigvec only works for square matrices
        if info[0][1].shape[0] != info[0][1].shape[1]:
            continue

        # Get the eigendecomposition of the weight matrix
        eigvals, eigvecs = [], []
        warnings.filterwarnings("ignore")  # converting to float from complex gives warning but is fine
        for _, p in info:
            eigval, eigvec = torch.eig(p)
            eigval, eigvec = eigval.to(torch.float), eigvec.to(torch.float)  # im is 0 anyway
            eigvals.append(eigval)
            eigvecs.append(eigvec)

        warnings.resetwarnings()  # now, warnings are useful again

        eigvals = [eigval.abs().mean() for eigval in eigvals]
        eigval_diff = (
            (torch.max(torch.tensor(eigvals)) - torch.min(torch.tensor(eigvals)))
            / torch.mean(torch.tensor(eigvals))
        )
        eigval_diffs.append(eigval_diff)

        eigvec_angle = torch.mean(
            torch.tensor(
                [vec_angle(v1, v2) for vecs in eigvecs for v1, v2 in itertools.combinations(vecs, 2)]
            )
        )
        eigvec_angles.append(eigvec_angle)

    return (
        torch.mean(torch.tensor(abs_mean_vals)),
        torch.mean(torch.tensor(abs_mean_diffs)),
        torch.mean(torch.tensor(eigval_diffs)),
        torch.mean(torch.tensor(eigvec_angles))
    )
//  [python]
def get_weight_infos(models: list[MLP]) -> tuple[torch.Tensor, ...]:
    abs_mean_vals = []
    abs_mean_diffs = []
    eigval_diffs = []
    eigvec_angles = []
    for info in zip(*[model.named_parameters() for model in models]):
        if "weight" not in info[0][0] or info[0][1].ndim != 2:
            continue

        abs_means = torch.tensor([p.abs().mean() for _, p in info])
        abs_mean_vals.append(torch.mean(abs_means))
        abs_mean_diff = (
                (torch.max(abs_means) - torch.min(abs_means))
                / torch.mean(abs_means)  # always positive due to squaring
        )
        abs_mean_diffs.append(abs_mean_diff)

        # Eigval & eigvec only works for square matrices
        if info[0][1].shape[0] != info[0][1].shape[1]:
            continue

        # Get the eigendecomposition of the weight matrix
        eigvals, eigvecs = [], []
        warnings.filterwarnings("ignore")  # converting to float from complex gives warning but is fine
        for _, p in info:
            eigval, eigvec = torch.linalg.eig(p)
            eigval, eigvec = eigval.to(torch.float), eigvec.to(torch.float)  # im is 0 anyway
            eigvals.append(eigval)
            eigvecs.append(eigvec)

        warnings.resetwarnings()  # now, warnings are useful again

        eigvals = [eigval.abs().mean() for eigval in eigvals]
        eigval_diff = (
            (torch.max(torch.tensor(eigvals)) - torch.min(torch.tensor(eigvals)))
            / torch.mean(torch.tensor(eigvals))
        )
        eigval_diffs.append(eigval_diff)

        eigvec_angle = torch.mean(
            torch.tensor(
                [vec_angle(v1, v2) for vecs in eigvecs for v1, v2 in itertools.combinations(vecs, 2)]
            )
        )
        eigvec_angles.append(eigvec_angle)

    return (
        torch.mean(torch.tensor(abs_mean_vals)),
        torch.mean(torch.tensor(abs_mean_diffs)),
        torch.mean(torch.tensor(eigval_diffs)),
        torch.mean(torch.tensor(eigvec_angles))
    )
//  [python]
def compute_diversity_loss(phi_sub, phi_input):
    def compute_diversity(phi):
        phi = f.normalize(phi, p=2, dim=1)
        S_B = torch.mm(phi, phi.t())
        eig_vals, eig_vecs = torch.symeig(S_B, eigenvectors=True)
        return eig_vals.float(), eig_vecs.float()

    def normalize_min_max(eig_vals):
        min_v, max_v = torch.min(eig_vals), torch.max(eig_vals)
        return (eig_vals - min_v + 1e-9) / (max_v - min_v + 1e-9)

    sub_eig_vals, sub_eig_vecs = compute_diversity(phi_sub)
    input_eig_vals, input_eig_vecs = compute_diversity(phi_input)
    magnitude_loss = torch.exp(-torch.norm(input_eig_vals-sub_eig_vals, p=2, dim=0, keepdim=True))
    weighted_structure_loss = torch.sum(torch.norm(torch.mul(sub_eig_vecs, input_eig_vecs), p=2, dim=1, keepdim=True))

    return magnitude_loss + weighted_structure_loss
//  [python]
def compute_diversity_loss(phi_sub, phi_input):
    def compute_diversity(phi):
        phi = f.normalize(phi, p=2, dim=1)
        S_B = torch.mm(phi, phi.t())
        eig_vals, eig_vecs = torch.linalg.eig(S_B)
        return eig_vals.float(), eig_vecs.float()

    def normalize_min_max(eig_vals):
        min_v, max_v = torch.min(eig_vals), torch.max(eig_vals)
        return (eig_vals - min_v + 1e-9) / (max_v - min_v + 1e-9)

    sub_eig_vals, sub_eig_vecs = compute_diversity(phi_sub)
    input_eig_vals, input_eig_vecs = compute_diversity(phi_input)
    magnitude_loss = torch.exp(-torch.norm(input_eig_vals-sub_eig_vals, p=2, dim=0, keepdim=True))
    weighted_structure_loss = torch.sum(torch.norm(torch.mul(sub_eig_vecs, input_eig_vecs), p=2, dim=1, keepdim=True))

    return magnitude_loss + weighted_structure_loss
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2] # Only used in testing
        return w, P, B, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.eig(AP) # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    B = contract('ij, j -> i', V_inv, B.to(V)) # V^* B
    P = contract('ij, ...j -> ...i', V_inv, P.to(V)) # V^* P


    return w, P, B, V
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2] # Only used in testing
        return w, P, B, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.linalg.eig(AP) # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    B = contract('ij, j -> i', V_inv, B.to(V)) # V^* B
    P = contract('ij, ...j -> ...i', V_inv, P.to(V)) # V^* P


    return w, P, B, V
//  [python]
def get_head_agg_weight(self, num_users, Vars, Hs, *args, **kwargs):
        device = Hs[0][0].device
        num_cls = Hs[0].shape[0]  # number of classes
        d = Hs[0].shape[1]  # dimension of feature representation
        avg_weight = []
        for i in range(num_users):
            # ---------------------------------------------------------------------------
            # variance ter
            v = torch.tensor(Vars, device=device)
            # ---------------------------------------------------------------------------
            # bias term
            h_ref = Hs[i]
            dist = torch.zeros((num_users, num_users), device=device)
            for j1, j2 in pairwise(tuple(range(num_users))):
                h_j1 = Hs[j1]
                h_j2 = Hs[j2]
                h = torch.zeros((d, d), device=device)
                for k in range(num_cls):
                    h += torch.mm((h_ref[k] - h_j1[k]).reshape(d, 1), (h_ref[k] - h_j2[k]).reshape(1, d))
                dj12 = torch.trace(h)
                dist[j1][j2] = dj12
                dist[j2][j1] = dj12
            # QP solver
            p_matrix = torch.diag(v) + dist
            p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
            evals, evecs = torch.eig(torch.tensor(p_matrix))
            # for numerical stablity
            p_matrix_new = 0
            p_matrix_new = 0
            for ii in range(num_users):
                if evals[ii].real >= 0.01:
                    p_matrix_new += evals[ii].real * torch.mm(evecs[:, ii].reshape(num_users, 1),
                                                              evecs[:, ii].reshape(1, num_users))
            p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix) >= 0.0) else p_matrix
            # solve QP
            alpha = 0
            eps = 1e-3
            if np.all(np.linalg.eigvals(p_matrix) >= 0):
                alphav = cvx.Variable(num_users)
                obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
                prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
                prob.solve()
                alpha = alphav.value
                alpha = [(i) * (i > eps) for i in alpha]  # zero-out small weights (<eps)
                if i == 0:
                    print('({}) Agg Weights of Classifier Head'.format(i + 1))
                    print(alpha, '\n')
            else:
                alpha = None  # if no solution for the optimization problem, use local classifier only

            avg_weight.append(alpha)

        return avg_weight
//  [python]
def get_head_agg_weight(self, num_users, Vars, Hs, *args, **kwargs):
        device = Hs[0][0].device
        num_cls = Hs[0].shape[0]  # number of classes
        d = Hs[0].shape[1]  # dimension of feature representation
        avg_weight = []
        for i in range(num_users):
            # ---------------------------------------------------------------------------
            # variance ter
            v = torch.tensor(Vars, device=device)
            # ---------------------------------------------------------------------------
            # bias term
            h_ref = Hs[i]
            dist = torch.zeros((num_users, num_users), device=device)
            for j1, j2 in pairwise(tuple(range(num_users))):
                h_j1 = Hs[j1]
                h_j2 = Hs[j2]
                h = torch.zeros((d, d), device=device)
                for k in range(num_cls):
                    h += torch.mm((h_ref[k] - h_j1[k]).reshape(d, 1), (h_ref[k] - h_j2[k]).reshape(1, d))
                dj12 = torch.trace(h)
                dist[j1][j2] = dj12
                dist[j2][j1] = dj12
            # QP solver
            p_matrix = torch.diag(v) + dist
            p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
            evals, evecs = torch.linalg.eig(torch.tensor(p_matrix))
            # for numerical stablity
            p_matrix_new = 0
            p_matrix_new = 0
            for ii in range(num_users):
                if evals[ii].real >= 0.01:
                    p_matrix_new += evals[ii].real * torch.mm(evecs[:, ii].reshape(num_users, 1),
                                                              evecs[:, ii].reshape(1, num_users))
            p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix) >= 0.0) else p_matrix
            # solve QP
            alpha = 0
            eps = 1e-3
            if np.all(np.linalg.eigvals(p_matrix) >= 0):
                alphav = cvx.Variable(num_users)
                obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
                prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
                prob.solve()
                alpha = alphav.value
                alpha = [(i) * (i > eps) for i in alpha]  # zero-out small weights (<eps)
                if i == 0:
                    print('({}) Agg Weights of Classifier Head'.format(i + 1))
                    print(alpha, '\n')
            else:
                alpha = None  # if no solution for the optimization problem, use local classifier only

            avg_weight.append(alpha)

        return avg_weight
//  [python]
def observe(self, x : pd.DataFrame, y : np.ndarray):
        """
        Observe new data
        """
        if self.cx is not None:
            assert (self.cx.values == x[self.cx.columns].values).all()
       
        if y.min() < self._best_y:
            self._best_y = y.min()
            self._best_x = x.iloc[[y.reshape(-1).argmin()]]

        self.n_eval += y.shape[0]

        y       = y.reshape(-1)
        lb      = self.space.opt_lb.view(-1).float()
        ub      = self.space.opt_ub.view(-1).float()
        cx, cxe = self.space.transform(x)
        cx      = torch.cat([cx, cxe.float()], dim = 1)

        self.px = x.iloc[y.argsort()[:self.pop_size]].copy()
        px, pxe = self.space.transform(self.px)
        px      = torch.cat([px, pxe.float()], dim = 1)

        mu_old        = self.mu.clone()
        self.mu       = (px.t() * self.weights).sum(axis = 1)

        # D, P          = self.C.eig(True)
        D, P       = torch.eig(self.C, eigenvectors=True)
        P          = P.real
        D          = D.real * torch.eye(self.dim)
        Dinv       = (1. / D.diag()) * torch.eye(self.dim)
        C_half_inv = P.mm(Dinv.sqrt()).mm(P.t())

        self.p_sigma  = (1 - self.cs) * self.p_sigma + np.sqrt(self.cs * (2 - self.cs) * self.mu_eff) * C_half_inv.mm((self.mu - mu_old).view(-1, 1) / self.sigma).view(self.p_sigma.shape)

        h_sig = 0.
        gen   = self.n_eval / self.child_size
        if self.p_sigma.norm() / np.sqrt(1 - (1 - self.cs)**(2 * gen + 1)) < (1.4 + 2 / (self.dim + 1)) * self.norm_rand:
            h_sig = 1.0
        self.p_c = (1 - self.cc)  * self.p_c + h_sig * np.sqrt(self.cc  * (2 - self.cc)   * self.mu_eff) * ((self.mu - mu_old) / self.sigma)

        self.sigma *= np.exp((self.cs / self.damp) * (self.p_sigma.norm() / self.norm_rand - 1) )
        if not torch.isfinite(self.sigma):
            self.sigma   = self.init_sigma()
            self.p_sigma = torch.zeros(self.dim)
        self.sigma = torch.min(2 * (ub - lb).norm(), self.sigma)

        # rank-mu estimation
        C_mu = torch.zeros(self.dim, self.dim)
        for i in range(self.pop_size):
            y     = (px[i] - mu_old).view(-1, 1) / self.sigma
            C_mu += self.weights[i] * y.mm(y.t())
        
        # rank-1 estimation
        C_r1 = self.p_c.view(-1, 1).mm(self.p_c.view(1, -1)) + (1 - h_sig) * self.cc * (2 - self.cc) * self.C

        self.C = (1 - self.c_r1 - self.c_mu) * self.C + self.c_r1 * C_r1 + self.c_mu * C_mu