jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def ss_update(self,SExx,SEyx,SEyy,n,lr=1.0):

        if self.X_mask is not None:
            SExx = SExx*self.X_mask*self.X_mask.transpose(-2,-1)
            SEyx = SEyx*self.X_mask
            invV = self.invV_0 + SExx
            muinvV = self.mu_0@self.invV_0 + SEyx
            mu = muinvV@invV.inverse()
            mu = mu*self.X_mask
        else:
            invV = self.invV_0 + SExx
            muinvV = self.mu_0@self.invV_0 + SEyx
            mu = torch.solve(muinvV.transpose(-2,-1),invV)[0].transpose(-2,-1)
            # mu = muinvV@invV.inverse()

        if self.mask is not None:  # Assumes mask is same for whole batch
            V = invV.inverse()
            U = self.invU.EinvSigma().inverse()
            Astar = V.unsqueeze(-3).unsqueeze(-2)*U.unsqueeze(-2).unsqueeze(-1)
            A  = Astar[...,~self.mask,:,:][...,:,~self.mask]            
            b = mu[...,~self.mask]
            gamma = torch.zeros_like(mu)
            gamma[...,~self.mask] = torch.solve(b,A)[0]
            mu = mu - U@gamma@V
            mu = mu*self.mask

        SEyy = SEyy - mu@invV@mu.transpose(-2,-1)
        SEyy = SEyy + self.mu_0@self.invV_0@self.mu_0.transpose(-2,-1)
        self.invU.ss_update(SEyy.diagonal(dim1=-2,dim2=-1), n.unsqueeze(-1), lr)
        self.invV = (invV-self.invV)*lr + self.invV
        self.invV = 0.5*(self.invV + self.invV.transpose(-2,-1))
        self.mu = (mu-self.mu)*lr + self.mu
        if(self.mask is not None):
            self.mu = self.mu*self.mask

#        self.invV_d, self.invV_v = torch.linalg.eigh(self.invV) 
#        self.V = self.invV_v@(1.0/self.invV_d.unsqueeze(-1)*self.invV_v.transpose(-2,-1))
#        self.logdetinvV = self.invV_d.log().sum(-1)
        self.V = self.invV.inverse()
        self.logdetinvV = self.invV.logdet()

        if self.X_mask is not None:
#            self.V = self.V * self.X_mask * self.X_mask.transpose(-2,-1)
#            self.invV = self.invV * self.X_mask * self.X_mask.transpose(-2,-1)
            self.mu = self.mu*self.X_mask
#            self.logdetinvV = self.logdetinvV - self.logdetinvV_0*(~self.X_mask).sum(-1).sum(-1)

        if self.uniform_precision==True:
            self.invU.gamma.alpha = self.invU.gamma.alpha.sum(-1,keepdim=True)
//  [python]
def ss_update(self,SExx,SEyx,SEyy,n,lr=1.0):

        if self.X_mask is not None:
            SExx = SExx*self.X_mask*self.X_mask.transpose(-2,-1)
            SEyx = SEyx*self.X_mask
            invV = self.invV_0 + SExx
            muinvV = self.mu_0@self.invV_0 + SEyx
            mu = muinvV@invV.inverse()
            mu = mu*self.X_mask
        else:
            invV = self.invV_0 + SExx
            muinvV = self.mu_0@self.invV_0 + SEyx
            mu = torch.linalg.solve(invV,muinvV.transpose(-2,-1)).transpose(-2,-1)
            # mu = muinvV@invV.inverse()

        if self.mask is not None:  # Assumes mask is same for whole batch
            V = invV.inverse()
            U = self.invU.EinvSigma().inverse()
            Astar = V.unsqueeze(-3).unsqueeze(-2)*U.unsqueeze(-2).unsqueeze(-1)
            A  = Astar[...,~self.mask,:,:][...,:,~self.mask]            
            b = mu[...,~self.mask]
            gamma = torch.zeros_like(mu)
            gamma[...,~self.mask] = torch.linalg.solve(A,b)
            mu = mu - U@gamma@V
            mu = mu*self.mask

        SEyy = SEyy - mu@invV@mu.transpose(-2,-1)
        SEyy = SEyy + self.mu_0@self.invV_0@self.mu_0.transpose(-2,-1)
        self.invU.ss_update(SEyy.diagonal(dim1=-2,dim2=-1), n.unsqueeze(-1), lr)
        self.invV = (invV-self.invV)*lr + self.invV
        self.invV = 0.5*(self.invV + self.invV.transpose(-2,-1))
        self.mu = (mu-self.mu)*lr + self.mu
        if(self.mask is not None):
            self.mu = self.mu*self.mask

#        self.invV_d, self.invV_v = torch.linalg.eigh(self.invV) 
#        self.V = self.invV_v@(1.0/self.invV_d.unsqueeze(-1)*self.invV_v.transpose(-2,-1))
#        self.logdetinvV = self.invV_d.log().sum(-1)
        self.V = self.invV.inverse()
        self.logdetinvV = self.invV.logdet()

        if self.X_mask is not None:
#            self.V = self.V * self.X_mask * self.X_mask.transpose(-2,-1)
#            self.invV = self.invV * self.X_mask * self.X_mask.transpose(-2,-1)
            self.mu = self.mu*self.X_mask
#            self.logdetinvV = self.logdetinvV - self.logdetinvV_0*(~self.X_mask).sum(-1).sum(-1)

        if self.uniform_precision==True:
            self.invU.gamma.alpha = self.invU.gamma.alpha.sum(-1,keepdim=True)
//  [python]
def pca_torch(X):
    # TODO 2Change this to do SVD, because it is stable and computationally
    # less intensive.
    covariance = torch.transpose(X, 1, 0) @ X
    S, U = torch.eig(covariance, eigenvectors=True)
    return S, U
//  [python]
def pca_torch(X):
    covariance = torch.transpose(X, 1, 0) @ X
    S, U = torch.linalg.eig(covariance)
    return S, U
//  [python]
def test_serialization_parametrization(self):
        r"""Test that it is possible to serialize a parametrized model via state_dict"""

        # A stateful parametrization
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("id", torch.eye(n))
                self.register_buffer("B", torch.empty(n, n))
                init.orthogonal_(self.B)

            def forward(self, X):
                A = X.triu(1)
                A = A - A.T
                return self.B @ torch.solve(self.id - A, self.id + A)[0]

        def get_model():
            model = torch.nn.Sequential(
                torch.nn.Linear(5, 5),
                torch.nn.ReLU(),
                torch.nn.Linear(5, 1),
            )

            parametrize.register_parametrization(model[0], "weight", Orthogonal(5))
            return model

        model = get_model()

        prev_weight = model[0].weight
        prev_B = model[0].parametrizations.weight[0].B

        new_model = get_model()
        with TemporaryFileName() as fname:
            torch.save(model.state_dict(), fname)
            new_model.load_state_dict(torch.load(fname))

        # Integrity tests
        self.assertTrue(parametrize.is_parametrized(new_model[0], "weight"))
        self.assertEqual(prev_weight, new_model[0].weight)
        self.assertEqual(prev_B, new_model[0].parametrizations.weight[0].B)

        # Trying to save the whole parametrized model raises
        with self.assertRaisesRegex(RuntimeError, "state_dict"):
            with TemporaryFileName() as fname:
                torch.save(model, fname)
//  [python]
def test_serialization_parametrization(self):
        r"""Test that it is possible to serialize a parametrized model via state_dict"""

        # A stateful parametrization
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("id", torch.eye(n))
                self.register_buffer("B", torch.empty(n, n))
                init.orthogonal_(self.B)

            def forward(self, X):
                A = X.triu(1)
                A = A - A.T
                return self.B @ torch.linalg.solve(self.id + A, self.id - A)

        def get_model():
            model = torch.nn.Sequential(
                torch.nn.Linear(5, 5),
                torch.nn.ReLU(),
                torch.nn.Linear(5, 1),
            )

            parametrize.register_parametrization(model[0], "weight", Orthogonal(5))
            return model

        model = get_model()

        prev_weight = model[0].weight
        prev_B = model[0].parametrizations.weight[0].B

        new_model = get_model()
        with TemporaryFileName() as fname:
            torch.save(model.state_dict(), fname)
            new_model.load_state_dict(torch.load(fname))

        # Integrity tests
        self.assertTrue(parametrize.is_parametrized(new_model[0], "weight"))
        self.assertEqual(prev_weight, new_model[0].weight)
        self.assertEqual(prev_B, new_model[0].parametrizations.weight[0].B)

        # Trying to save the whole parametrized model raises
        with self.assertRaisesRegex(RuntimeError, "state_dict"):
            with TemporaryFileName() as fname:
                torch.save(model, fname)
//  [python]
def test_caching_parametrization(self):
        r"""Test the caching system of a parametrization"""

        # Define a couple matrix parametrizations
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.solve(Id - X, Id + X)[0]

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        # Test that the caching system works
        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))
//  [python]
def test_caching_parametrization(self):
        r"""Test the caching system of a parametrization"""

        # Define a couple matrix parametrizations
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.linalg.solve(Id + X, Id - X)

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        # Test that the caching system works
        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))
//  [python]
def forward(self, X, Y):

        n, k = X.shape[:2]  # n = 77, k =2
        device = X.device

        Z = torch.zeros(1, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device) # [1, 80, 80]
        K = K_matrix(X, X)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        Q = torch.solve(Z, L)[0]
        # Q = torch.linalg.solve(L, Z)
        return Q[:, :k], Q[:, k:]
//  [python]
def forward(self, X, Y):

        n, k = X.shape[:2]  # n = 77, k =2
        device = X.device

        Z = torch.zeros(1, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device) # [1, 80, 80]
        K = K_matrix(X, X)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        # Q = torch.solve(Z, L)[0]
        Q = torch.linalg.solve(L, Z)
        return Q[:, :k], Q[:, k:]
//  [python]
def test_caching_parametrization_with_transfer_parametrizations_and_params(self):
        r"""Test that transferring parametrizations doesn't cause issues with caching"""

        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.solve(Id - X, Id + X)[0]

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        to_model = nn.Linear(5, 5)
        parametrize.transfer_parametrizations_and_params(model, to_model)

        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))

            A = to_model.weight
            B = to_model.weight
            self.assertEqual(id(A), id(B))

            # test that the results are distinct objects for each module
            self.assertNotEqual(id(A), id(X))
//  [python]
def test_caching_parametrization_with_transfer_parametrizations_and_params(self):
        r"""Test that transferring parametrizations doesn't cause issues with caching"""

        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.linalg.solve(Id + X, Id - X)

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        to_model = nn.Linear(5, 5)
        parametrize.transfer_parametrizations_and_params(model, to_model)

        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))

            A = to_model.weight
            B = to_model.weight
            self.assertEqual(id(A), id(B))

            # test that the results are distinct objects for each module
            self.assertNotEqual(id(A), id(X))
//  [python]
def logdet(x):
    """

    Args:
        x: 2D positive semidefinite matrix.

    Returns: log determinant of x

    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print((torch.log(torch.eig(x.data)[0])))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def logdet(x):
    """

    Args:
        x: 2D positive semidefinite matrix.

    Returns: log determinant of x

    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print((torch.log(torch.linalg.eig(x.data)[0])))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def pca_torch(X, no_dims=50, device=None):
    print("Preprocessing the data using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.eig(torch.mm(X.t(), X), True)
    # split M real
    for i in range(d):
        if l[i, 1] != 0:
            M[:, i+1] = M[:, i]
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def pca_torch(X, no_dims=50, device=None):
    print("Preprocessing the data using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.linalg.eig(torch.mm(X.t(), X))
    # split M real
    for i in range(d):
        if l[i, 1] != 0:
            M[:, i+1] = M[:, i]
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def test_initialization_parametrization(self):
        r"""Test that it is possible to initialize a parametrization when it
        implements a `right_inverse` method
        """

        class Skew(nn.Module):
            def forward(self, X):
                A = X.triu(1)
                return A - A.T

            def is_skew(self, A):
                return torch.allclose(A, -A.T, atol=1e-6)

            def right_inverse(self, X):
                if not self.is_skew(X):
                    raise ValueError("The matrix is not skew-symmetric.")
                return X.triu(1)

        # Implements a Cayley map where right_inverse is not quite the inverse of forward
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("B", torch.eye(n))

            def forward(self, X):
                Id = torch.eye(X.size(0))
                return self.B @ torch.solve(Id - X, Id + X)[0]

            def is_orthogonal(self, X):
                Id = torch.eye(X.size(0))
                return torch.allclose(X.T @ X, Id, atol=1e-4)

            def right_inverse(self, X):
                if not self.is_orthogonal(X):
                    raise ValueError("The input is not orthogonal.")
                # cayley(0) == Id, so B @ cayley(0) == B
                self.B = X
                return torch.zeros_like(X)

        N = 5
        model = nn.Linear(N, N)
        # Register the skew-symmetric constraint. The result is now skew-symmetric
        skew = Skew()
        # Make the weight skew-symmetric before registering the parametrization
        with torch.no_grad():
            model.weight.set_(skew(model.weight))
        parametrize.register_parametrization(model, "weight", skew)
        X = torch.rand(N, N)
        # X is not skew-symmetric, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        # Make X skew-symmetric
        X = X - X.T
        model.weight = X
        self.assertEqual(model.parametrizations.weight.original, X.triu(1))
        self.assertEqual(model.weight, X)

        # Having several parametrizations registered should work in the same way
        parametrize.register_parametrization(model, "weight", Orthogonal(N))
        # Register now the Cayley map. The result is now orthogonal
        X = torch.rand(N, N)
        # X is not orthogonal, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        init.orthogonal_(X)
        model.weight = X
        self.assertEqual(model.weight, X)
        self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X))
//  [python]
def test_initialization_parametrization(self):
        r"""Test that it is possible to initialize a parametrization when it
        implements a `right_inverse` method
        """

        class Skew(nn.Module):
            def forward(self, X):
                A = X.triu(1)
                return A - A.T

            def is_skew(self, A):
                return torch.allclose(A, -A.T, atol=1e-6)

            def right_inverse(self, X):
                if not self.is_skew(X):
                    raise ValueError("The matrix is not skew-symmetric.")
                return X.triu(1)

        # Implements a Cayley map where right_inverse is not quite the inverse of forward
        class Orthogonal(nn.Module):
            def __init__(self, n):
                super().__init__()
                self.register_buffer("B", torch.eye(n))

            def forward(self, X):
                Id = torch.eye(X.size(0))
                return self.B @ torch.linalg.solve(Id + X, Id - X)

            def is_orthogonal(self, X):
                Id = torch.eye(X.size(0))
                return torch.allclose(X.T @ X, Id, atol=1e-4)

            def right_inverse(self, X):
                if not self.is_orthogonal(X):
                    raise ValueError("The input is not orthogonal.")
                # cayley(0) == Id, so B @ cayley(0) == B
                self.B = X
                return torch.zeros_like(X)

        N = 5
        model = nn.Linear(N, N)
        # Register the skew-symmetric constraint. The result is now skew-symmetric
        skew = Skew()
        # Make the weight skew-symmetric before registering the parametrization
        with torch.no_grad():
            model.weight.set_(skew(model.weight))
        parametrize.register_parametrization(model, "weight", skew)
        X = torch.rand(N, N)
        # X is not skew-symmetric, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        # Make X skew-symmetric
        X = X - X.T
        model.weight = X
        self.assertEqual(model.parametrizations.weight.original, X.triu(1))
        self.assertEqual(model.weight, X)

        # Having several parametrizations registered should work in the same way
        parametrize.register_parametrization(model, "weight", Orthogonal(N))
        # Register now the Cayley map. The result is now orthogonal
        X = torch.rand(N, N)
        # X is not orthogonal, so it throws an error
        with self.assertRaises(ValueError):
            model.weight = X
        init.orthogonal_(X)
        model.weight = X
        self.assertEqual(model.weight, X)
        self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X))
//  [python]
def is_psd(mat):
    eigs = torch.sort(torch.eig(mat)[0][:, 0])[0]
    print(f'Min eigenvalue is {eigs[0]:.2e}')
    print(f'Max eigenvalue is {eigs[-1]:.2e}')
    if bool((mat == mat.T).all() and (eigs >= 0).all()):
        print('matrix is positive semi-definite')
    else:
        raise ValueError('Negative definite matrix from Chebnet')
//  [python]
def is_psd(mat):
    eigs = torch.sort(torch.linalg.eig(mat)[0][:, 0])[0]
    print(f'Min eigenvalue is {eigs[0]:.2e}')
    print(f'Max eigenvalue is {eigs[-1]:.2e}')
    if bool((mat == mat.T).all() and (eigs >= 0).all()):
        print('matrix is positive semi-definite')
    else:
        raise ValueError('Negative definite matrix from Chebnet')
//  [python]
def rbf(self, X):
        Xk = None
        nfeature    = X.shape[1]
        Degree      = self.degree
        Gamma       = self.gamma
        NComp       = self.n_component
        if Gamma is None:
            Gamma = 1 / nfeature
        if NComp is None:
            NComp = nfeature
        K = Gamma*(torch.mm(X, X.t()) + 1)**Degree
        ninstance = K.shape[0]
        eigvals, eigvecs = torch.eig(K, eigenvectors=True)
        for i in range(1, NComp + 1):
           Xk = torch.reshape(eigvecs[:, -i], (1, ninstance)) if Xk is None else torch.cat((Xk, torch.reshape(eigvecs[:, -i], (1, ninstance))))
        return Xk.t()
//  [python]
def rbf(self, X):
        Xk = None
        nfeature    = X.shape[1]
        Degree      = self.degree
        Gamma       = self.gamma
        NComp       = self.n_component
        if Gamma is None:
            Gamma = 1 / nfeature
        if NComp is None:
            NComp = nfeature
        K = Gamma*(torch.mm(X, X.t()) + 1)**Degree
        ninstance = K.shape[0]
        eigvals, eigvecs = torch.linalg.eig(K)
        for i in range(1, NComp + 1):
           Xk = torch.reshape(eigvecs[:, -i], (1, ninstance)) if Xk is None else torch.cat((Xk, torch.reshape(eigvecs[:, -i], (1, ninstance))))
        return Xk.t()
//  [python]
def transform(self, X, y, n_components):
        """
        And Why Inverse, In matrices, there is no concepts of division, thus multiplying with inverse
        matrix helps in acheiving what division does.
        :param X:
        :param y:
        :param n_components: Transforming from high dimension data to lower dimension n_components.
        :return: Transformed set of low dimensional X matrix
        """
        sw, sb = self.scatter_matrix(X, y)
        A = torch.mm(torch.pinverse(sw), sb)
        eigen_values, eigen_vectors = torch.eig(A, eigenvectors=True)
        eigen_sorted_index = torch.argsort(eigen_values[:, 0], descending=True)
        eigen_vectors_sorted = eigen_vectors[:, eigen_sorted_index]
        component_vector = eigen_vectors_sorted[:, 0:n_components]
        component_vector = component_vector.type(torch.DoubleTensor)
        transformed = torch.mm(X, component_vector)
        return transformed
//  [python]
def transform(self, X, y, n_components):
        """
        And Why Inverse, In matrices, there is no concepts of division, thus multiplying with inverse
        matrix helps in acheiving what division does.
        :param X:
        :param y:
        :param n_components: Transforming from high dimension data to lower dimension n_components.
        :return: Transformed set of low dimensional X matrix
        """
        sw, sb = self.scatter_matrix(X, y)
        A = torch.mm(torch.pinverse(sw), sb)
        eigen_values, eigen_vectors = torch.linalg.eig(A)
        eigen_sorted_index = torch.argsort(eigen_values[:, 0], descending=True)
        eigen_vectors_sorted = eigen_vectors[:, eigen_sorted_index]
        component_vector = eigen_vectors_sorted[:, 0:n_components]
        component_vector = component_vector.type(torch.DoubleTensor)
        transformed = torch.mm(X, component_vector)
        return transformed
//  [python]
def logdet(x):
    """

    Args:
        x: 2D positive semidefinite matrix.

    Returns: log determinant of x

    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print(torch.log(torch.eig(x.data)[0]))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def logdet(x):
    """

    Args:
        x: 2D positive semidefinite matrix.

    Returns: log determinant of x

    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print(torch.log(torch.linalg.eig(x.data)[0]))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def rot_evecs(M, theta):
    " Rotate eigenvectors of matrix M "
    evals, evecs = torch.eig(M, eigenvectors = True)
    evecs_rot = rot(evecs, theta)
    return spectrally_prescribed_matrix(evals, evecs_rot)
//  [python]
def rot_evecs(M, theta):
    " Rotate eigenvectors of matrix M "
    evals, evecs = torch.linalg.eig(M)
    evecs_rot = rot(evecs, theta)
    return spectrally_prescribed_matrix(evals, evecs_rot)
//  [python]
def get_largest_eigenvalue(X):
    eigs = torch.eig(X, eigenvectors=False).eigenvalues
    max_eign = eigs.max(dim=0)
    return max_eign.values[0]
//  [python]
def get_largest_eigenvalue(X):
    eigs = torch.linalg.eig(X).eigenvalues
    max_eign = eigs.max(dim=0)
    return max_eign.values[0]
//  [python]
def torch_sqrtm(m):
    """
    Square root of a matrix
    inspired by: https://discuss.pytorch.org/t/raising-a-tensor-to-a-fractional-power/93655/3
    """
    assert is_psd(m), 'Input matrix is not positive semi-definitive'

    evals, evecs = torch.eig(m, eigenvectors=True)  # get eigendecomposition
    evals = torch.relu(evals[:, 0])  # get real part of (real) eigenvalues
    evpow = evals ** (1 / 2)  # raise eigenvalues to fractional power

    # build exponentiated matrix from exponentiated eigenvalues
    mpow = torch.matmul(evecs, torch.matmul(torch.diag(evpow), torch.inverse(evecs)))
    return mpow
//  [python]
def torch_sqrtm(m):
    """
    Square root of a matrix
    inspired by: https://discuss.pytorch.org/t/raising-a-tensor-to-a-fractional-power/93655/3
    """
    assert is_psd(m), 'Input matrix is not positive semi-definitive'

    evals, evecs = torch.linalg.eig(m)  # get eigendecomposition
    evals = torch.relu(evals[:, 0])  # get real part of (real) eigenvalues
    evpow = evals ** (1 / 2)  # raise eigenvalues to fractional power

    # build exponentiated matrix from exponentiated eigenvalues
    mpow = torch.matmul(evecs, torch.matmul(torch.diag(evpow), torch.inverse(evecs)))
    return mpow
//  [python]
def NTK(pred):
    paranum = sum(p.numel() for p in model_h.parameters())
    A = torch.zeros((len(pred), paranum)).cuda()
    for index,pred_e in enumerate(pred):
        grad_e = grad(pred_e, model_h.parameters(),retain_graph=True,  create_graph=True) # 获得每一个预测对参数的梯度
        grad_one = torch.tensor([]).cuda() # 创建一个空的tensor
        for i in grad_e:
            grad_one = torch.cat((grad_one, i.flatten())) # 将每一个预测对参数的梯度预测都放入grad_one 中
        A[index] = grad_one # 存入A中，A的行是样本个数，列是梯度个数
    K = torch.mm(A, A.T) # 矩阵的乘法组成K矩阵，这个K矩阵的获得是利用了矩阵的乘法，矩阵的乘法可以加速获得K矩阵，这里要避免使用for循环
    L,_ = torch.eig(K.data.cpu()) # torch不能对gpu求特征值
    eig = torch.norm(L, dim=1)
    return eig
//  [python]
def NTK(pred):
    paranum = sum(p.numel() for p in model_h.parameters())
    A = torch.zeros((len(pred), paranum)).cuda()
    for index,pred_e in enumerate(pred):
        grad_e = grad(pred_e, model_h.parameters(),retain_graph=True,  create_graph=True) # 获得每一个预测对参数的梯度
        grad_one = torch.tensor([]).cuda() # 创建一个空的tensor
        for i in grad_e:
            grad_one = torch.cat((grad_one, i.flatten())) # 将每一个预测对参数的梯度预测都放入grad_one 中
        A[index] = grad_one # 存入A中，A的行是样本个数，列是梯度个数
    K = torch.mm(A, A.T) # 矩阵的乘法组成K矩阵，这个K矩阵的获得是利用了矩阵的乘法，矩阵的乘法可以加速获得K矩阵，这里要避免使用for循环
    L,_ = torch.linalg.eig(K.data.cpu()) # torch不能对gpu求特征值
    eig = torch.norm(L, dim=1)
    return eig
//  [python]
def get_exact_laplacian(states, r_states, n_actions=4):
    # Build adjacency matrix
    A = th.zeros((len(states), len(states)))
    states_map = r_states > -1
    cur_pos = r_states[states_map]
    left_pos = r_states[states_map.roll(shifts=-1, dims=1)]
    left_pos[left_pos == -1] = cur_pos[left_pos == -1]
    A[cur_pos, left_pos] = 1
    right_pos = r_states[states_map.roll(shifts=1, dims=1)]
    right_pos[right_pos == -1] = cur_pos[right_pos == -1]
    A[cur_pos, right_pos] = 1
    up_pos = r_states[states_map.roll(shifts=-1, dims=0)]
    up_pos[up_pos == -1] = cur_pos[up_pos == -1]
    A[cur_pos, up_pos] = 1
    down_pos = r_states[states_map.roll(shifts=1, dims=0)]
    down_pos[down_pos == -1] = cur_pos[down_pos == -1]
    A[cur_pos, down_pos] = 1

    # Build transition matrix
    P = A / n_actions
    P[range(len(states)), range(len(states))] += 1 - P.sum(axis=0)

    # Compute graph Laplacian
    D = A.sum(axis=-1)
    L = th.diag(D) - A

    # Eigendecompose the Laplacian
    eigenvalues, eigenvectors = th.eig(L, eigenvectors=True)
    idx = eigenvalues[:, 0].argsort()
    eigenvalues = eigenvalues[idx, 0]
    eigenvectors = eigenvectors[:, idx]

    return eigenvectors, eigenvalues, P
//  [python]
def get_exact_laplacian(states, r_states, n_actions=4):
    # Build adjacency matrix
    A = th.zeros((len(states), len(states)))
    states_map = r_states > -1
    cur_pos = r_states[states_map]
    left_pos = r_states[states_map.roll(shifts=-1, dims=1)]
    left_pos[left_pos == -1] = cur_pos[left_pos == -1]
    A[cur_pos, left_pos] = 1
    right_pos = r_states[states_map.roll(shifts=1, dims=1)]
    right_pos[right_pos == -1] = cur_pos[right_pos == -1]
    A[cur_pos, right_pos] = 1
    up_pos = r_states[states_map.roll(shifts=-1, dims=0)]
    up_pos[up_pos == -1] = cur_pos[up_pos == -1]
    A[cur_pos, up_pos] = 1
    down_pos = r_states[states_map.roll(shifts=1, dims=0)]
    down_pos[down_pos == -1] = cur_pos[down_pos == -1]
    A[cur_pos, down_pos] = 1

    # Build transition matrix
    P = A / n_actions
    P[range(len(states)), range(len(states))] += 1 - P.sum(axis=0)

    # Compute graph Laplacian
    D = A.sum(axis=-1)
    L = th.diag(D) - A

    # Eigendecompose the Laplacian
    eigenvalues, eigenvectors = th.linalg.eig(L)
    idx = eigenvalues[:, 0].argsort()
    eigenvalues = eigenvalues[idx, 0]
    eigenvectors = eigenvectors[:, idx]

    return eigenvectors, eigenvalues, P
//  [python]
def forward(self, features, A):
        b, n, d = features.shape
        assert (d == self.in_dim)
        if(len(A.shape) == 2):
            # A_norm = A + torch.eye(n).cuda()
            A_norm = A
            deg_mat = Comp_degree(A_norm)
            frac_degree = torch.FloatTensor(fractional_matrix_power(deg_mat.detach().cpu(),
                                                                    -0.5)).cuda()
            Laplacian = deg_mat - A_norm
            Laplacian_norm = frac_degree.matmul(Laplacian.matmul(frac_degree))
            
            landa, U = torch.eig(Laplacian_norm,eigenvectors=True)
    
            repeated_U_t = U.t().repeat(b, 1, 1)
            repeated_U = U.repeat(b, 1, 1)
        else:
            repeated_U_t = []
            repeated_U = []
            for i in range(A.shape[0]):
                # A_norm = A[i] + torch.eye(n).cuda()
                A_norm = A[i]
                deg_mat = Comp_degree(A_norm)
                frac_degree = torch.FloatTensor(fractional_matrix_power(deg_mat.detach().cpu(),
                                                                        -0.5)).cuda()
                Laplacian = deg_mat - A_norm
                Laplacian_norm = frac_degree.matmul(Laplacian.matmul(frac_degree))

                landa, U = torch.eig(Laplacian_norm, eigenvectors=True)

                repeated_U_t.append(U.t().view(1, U.shape[0], U.shape[1]))
                repeated_U.append(U.view(1, U.shape[0], U.shape[1]))
            repeated_U_t = torch.cat(repeated_U_t)
            repeated_U = torch.cat(repeated_U)
                
        agg_feats = torch.bmm(repeated_U_t, features)

        #### Adding MLP to GCN
        out = self.MLP(agg_feats.view(-1, d)).view(b, -1, self.out_dim)
        out = torch.bmm(repeated_U, out)
        # out = self.batchnorm(out).view(b, -1, self.out_dim)

        return out
//  [python]
def forward(self, features, A):
        b, n, d = features.shape
        assert (d == self.in_dim)
        if(len(A.shape) == 2):
            # A_norm = A + torch.eye(n).cuda()
            A_norm = A
            deg_mat = Comp_degree(A_norm)
            frac_degree = torch.FloatTensor(fractional_matrix_power(deg_mat.detach().cpu(),
                                                                    -0.5)).cuda()
            Laplacian = deg_mat - A_norm
            Laplacian_norm = frac_degree.matmul(Laplacian.matmul(frac_degree))
            
            landa, U = torch.linalg.eig(Laplacian_norm)
    
            repeated_U_t = U.t().repeat(b, 1, 1)
            repeated_U = U.repeat(b, 1, 1)
        else:
            repeated_U_t = []
            repeated_U = []
            for i in range(A.shape[0]):
                # A_norm = A[i] + torch.eye(n).cuda()
                A_norm = A[i]
                deg_mat = Comp_degree(A_norm)
                frac_degree = torch.FloatTensor(fractional_matrix_power(deg_mat.detach().cpu(),
                                                                        -0.5)).cuda()
                Laplacian = deg_mat - A_norm
                Laplacian_norm = frac_degree.matmul(Laplacian.matmul(frac_degree))

                landa, U = torch.linalg.eig(Laplacian_norm)

                repeated_U_t.append(U.t().view(1, U.shape[0], U.shape[1]))
                repeated_U.append(U.view(1, U.shape[0], U.shape[1]))
            repeated_U_t = torch.cat(repeated_U_t)
            repeated_U = torch.cat(repeated_U)
                
        agg_feats = torch.bmm(repeated_U_t, features)

        #### Adding MLP to GCN
        out = self.MLP(agg_feats.view(-1, d)).view(b, -1, self.out_dim)
        out = torch.bmm(repeated_U, out)
        # out = self.batchnorm(out).view(b, -1, self.out_dim)

        return out
//  [python]
def test_eig_reuse(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')

        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')
//  [python]
def test_eig_reuse(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.linalg.eig(X, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')

        torch.linalg.eig(X, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')
//  [python]
def load_sam_mat(self, root_path):

        sam_mat_path = join(root_path, "results", "sam_mat", "sam_mat.txt")
        sam_outpath = join(root_path, "results", "sam")
        
        if os.path.isfile(sam_mat_path):
            print("Loaded pca feature matrix from: ", sam_mat_path)
            return torch.Tensor(np.loadtxt(sam_mat_path)).cuda()
        
        elif os.path.isdir(sam_outpath):

            def pca(X, n_components=3):
                n = X.shape[0]
                mean = torch.mean(X, axis=0)
                X = X - mean
                covariance_matrix = (1 / n) * torch.matmul(X.T, X).float()  # An old torch bug: matmul float32->float16, 
                eigenvalues, eigenvectors = torch.eig(covariance_matrix, eigenvectors=True)
                eigenvalues = torch.norm(eigenvalues, dim=1)
                idx = torch.argsort(-eigenvalues)
                eigenvectors = eigenvectors[:, idx]
                proj_mat = eigenvectors[:, 0:n_components]
                
                return proj_mat

            imgs_p = sorted(glob.glob(join(sam_outpath, "*.npy")))
            n_img = len(imgs_p)
            # --- Load all feature map --- #
            imgs = []
            for i in range(n_img):
                imgs.append(torch.tensor(np.load(imgs_p[i])))
            imgs = torch.stack(imgs, dim=0)
            imgs_flat = imgs.reshape(-1, imgs.shape[-1])
            rand_idx = torch.randint(0, int(imgs_flat.shape[0]), [100000])

            pixs = imgs_flat[rand_idx, :]
            sam_norm = pixs / (pixs.norm(dim=1, keepdim=True) + 1e-6)
            proj_mat_3 = pca(sam_norm, n_components=3)

            os.makedirs(dirname(sam_mat_path), exist_ok=True)
            np.savetxt(sam_mat_path, proj_mat_3.numpy())
            print("Saved pca feature matrix to: ", sam_mat_path)

            return proj_mat_3.cuda()
        
        else:
            return None
//  [python]
def load_sam_mat(self, root_path):

        sam_mat_path = join(root_path, "results", "sam_mat", "sam_mat.txt")
        sam_outpath = join(root_path, "results", "sam")
        
        if os.path.isfile(sam_mat_path):
            print("Loaded pca feature matrix from: ", sam_mat_path)
            return torch.Tensor(np.loadtxt(sam_mat_path)).cuda()
        
        elif os.path.isdir(sam_outpath):

            def pca(X, n_components=3):
                n = X.shape[0]
                mean = torch.mean(X, axis=0)
                X = X - mean
                covariance_matrix = (1 / n) * torch.matmul(X.T, X).float()  # An old torch bug: matmul float32->float16, 
                eigenvalues, eigenvectors = torch.linalg.eig(covariance_matrix)
                eigenvalues = torch.norm(eigenvalues, dim=1)
                idx = torch.argsort(-eigenvalues)
                eigenvectors = eigenvectors[:, idx]
                proj_mat = eigenvectors[:, 0:n_components]
                
                return proj_mat

            imgs_p = sorted(glob.glob(join(sam_outpath, "*.npy")))
            n_img = len(imgs_p)
            # --- Load all feature map --- #
            imgs = []
            for i in range(n_img):
                imgs.append(torch.tensor(np.load(imgs_p[i])))
            imgs = torch.stack(imgs, dim=0)
            imgs_flat = imgs.reshape(-1, imgs.shape[-1])
            rand_idx = torch.randint(0, int(imgs_flat.shape[0]), [100000])

            pixs = imgs_flat[rand_idx, :]
            sam_norm = pixs / (pixs.norm(dim=1, keepdim=True) + 1e-6)
            proj_mat_3 = pca(sam_norm, n_components=3)

            os.makedirs(dirname(sam_mat_path), exist_ok=True)
            np.savetxt(sam_mat_path, proj_mat_3.numpy())
            print("Saved pca feature matrix to: ", sam_mat_path)

            return proj_mat_3.cuda()
        
        else:
            return None
//  [python]
def test_eig_invalid_input(self, device, dtype):
        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'input should be 2 dimensional',
            lambda: torch.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should be square',
            lambda: torch.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.eig(np.nan * torch.ones((2, 2))))
//  [python]
def test_eig_invalid_input(self, device, dtype):
        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'input should be 2 dimensional',
            lambda: torch.linalg.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should be square',
            lambda: torch.linalg.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.nan * torch.ones((2, 2))))
//  [python]
def test_eig_non_contiguous(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
//  [python]
def test_eig_non_contiguous(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        e, v = torch.linalg.eig(X)
        Xhat = torch.mm(torch.mm(v, torch.diag_embed(e[..., 0])), v.t())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
//  [python]
def test_eig_basic(self, device, dtype):
        a = torch.tensor([[1.96, 0.00, 0.00, 0.00, 0.00],
                          [-6.49, 3.80, 0.00, 0.00, 0.00],
                          [-0.47, -6.39, 4.17, 0.00, 0.00],
                          [-7.20, 1.50, -1.51, 5.70, 0.00],
                          [-0.65, -6.34, 2.67, 1.80, -7.10]],
                         dtype=dtype, device=device).t()
        e = torch.eig(a)[0]
        ee, vv = torch.eig(a, True)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)
        #
        # compare with numpy
        np_e, np_v = np.linalg.eig(a.cpu().numpy())
        # np_e.shape == (n, 2), where each column contain the real and
        # imaginary parts of the result
        self.assertEqual(ee[:, 0], np_e)  # real part
        self.assertEqual(ee[:, 1], torch.zeros(ee.shape[0], dtype=dtype))  # imaginary part
        self.assertEqual(vv, np_v)
//  [python]
def test_eig_basic(self, device, dtype):
        a = torch.tensor([[1.96, 0.00, 0.00, 0.00, 0.00],
                          [-6.49, 3.80, 0.00, 0.00, 0.00],
                          [-0.47, -6.39, 4.17, 0.00, 0.00],
                          [-7.20, 1.50, -1.51, 5.70, 0.00],
                          [-0.65, -6.34, 2.67, 1.80, -7.10]],
                         dtype=dtype, device=device).t()
        e = torch.linalg.eig(a)[0]
        ee, vv = torch.linalg.eig(a, eigenvectors=True)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.linalg.eig(a, eigenvectors=True, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)
        #
        # compare with numpy
        np_e, np_v = np.linalg.eig(a.cpu().numpy())
        # np_e.shape == (n, 2), where each column contain the real and
        # imaginary parts of the result
        self.assertEqual(ee[:, 0], np_e)  # real part
        self.assertEqual(ee[:, 1], torch.zeros(ee.shape[0], dtype=dtype))  # imaginary part
        self.assertEqual(vv, np_v)
//  [python]
def max_eigen(self, method='power'):
        assert method in ['torch', 'power']
        if method == 'torch':
            # kinda time consuming, but accurate
            assert self.mat is not None
            eig, _ = torch.eig(self.mat)
            eig = eig[:, 0]
            max_eig = eig.max().item()
        else:
            # use power method
            assert self.smat is not None
            max_eig = self.smat.max_eigenvalue()
        return max_eig
//  [python]
def max_eigen(self, method='power'):
        assert method in ['torch', 'power']
        if method == 'torch':
            # kinda time consuming, but accurate
            assert self.mat is not None
            eig = torch.linalg.eig(self.mat)
            eig = eig[:, 0]
            max_eig = eig.max().item()
        else:
            # use power method
            assert self.smat is not None
            max_eig = self.smat.max_eigenvalue()
        return max_eig
//  [python]
def is_positive_definite(tensor):
	"""Bool check if 2D matrix is positive definite

	:type tensor: torch.Tensor"""

	eig, _ = torch.eig(tensor)
	eig = eig[:, 0]  # only real part
	return torch.all(eig > 0)
//  [python]
def is_positive_definite(tensor):
	"""Bool check if 2D matrix is positive definite

	:type tensor: torch.Tensor"""

	eig = torch.linalg.eig(tensor).real
	return torch.all(eig > 0)
//  [python]
def infer_dimension(self, update_max_rank=True, use_delta=True):
        delta = 0.
        if use_delta:
            delta = self.subspace.delta

        _, var, subspace = self.get_space()
        subspace /= (self.n_models.item() - 1) ** 0.5
        tr_sigma = var.sum()

        spectrum, _ = torch.eig(subspace @ subspace.t())
        spectrum, _ = torch.sort(spectrum[:, 0], descending=True)

        new_max_rank, ll, _ = _infer_dimension_(spectrum.numpy(), tr_sigma.numpy(),
                                                self.n_models.item(), self.num_parameters, delta)

        if new_max_rank + 1 == self.subspace.max_rank and update_max_rank:
            self.subspace.max_rank += 1
//  [python]
def infer_dimension(self, update_max_rank=True, use_delta=True):
        delta = 0.
        if use_delta:
            delta = self.subspace.delta

        _, var, subspace = self.get_space()
        subspace /= (self.n_models.item() - 1) ** 0.5
        tr_sigma = var.sum()

        spectrum, _ = torch.linalg.eig(subspace @ subspace.t())
        spectrum, _ = torch.sort(spectrum[:, 0], descending=True)

        new_max_rank, ll, _ = _infer_dimension_(spectrum.numpy(), tr_sigma.numpy(),
                                                self.n_models.item(), self.num_parameters, delta)

        if new_max_rank + 1 == self.subspace.max_rank and update_max_rank:
            self.subspace.max_rank += 1
//  [python]
def PCA(X, k=20):
        X = torch.from_numpy(X)
        X_mean = torch.mean(X, 0)
        X = X - X_mean.expand_as(X)

        v, w = torch.eig(torch.mm(X, torch.t(X)), eigenvectors=True)
        return torch.mm(w[:k, :], X).numpy()
//  [python]
def PCA(X, k=20):
        X = torch.from_numpy(X)
        X_mean = torch.mean(X, 0)
        X = X - X_mean.expand_as(X)

        v, w = torch.linalg.eig(torch.mm(X, torch.t(X)))
        return torch.mm(w[:, :k], X).numpy()