jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_SVDSYMARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D0, U0= torch.symeig(M)
    S0,p= torch.sort(torch.abs(D0),descending=True)

    U,S,V= SVDSYMARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
        < S0[0]*(m**2)*1e-14 )
//  [python]
def test_SVDSYMARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D0, U0= torch.linalg.eigh(M)
    S0,p= torch.sort(torch.abs(D0),descending=True)

    U,S,V= SVDSYMARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
        < S0[0]*(m**2)*1e-14 )
//  [python]
def forward(ctx, x):
        ws, vs = torch.symeig(x, eigenvectors=True)
        ctx.save_for_backward(ws, vs)
        return ws.pow(2).sum()
//  [python]
def forward(ctx, x):
        ws, vs = torch.linalg.eigh(x)
        ctx.save_for_backward(ws, vs)
        return ws.pow(2).sum()
//  [python]
def test_hessian_eigs_positive_lanczos(self):
        "Compare `eff_dim.hessian_eigs_positive_lanczos` with hessian then `torch.symeig`"
        model, fn, loader = self.example_model_loss_data()
        M = eff_dim.hessian(model.parameters(), fn, loader)
        ref_vals, ref_vecs = torch.symeig(M, eigenvectors=True)
        num_negative = (ref_vals < 0).sum().item()

        # All values/vectors:
        vals, vecs = eff_dim.hessian_eigs_positive_lanczos(
            model.parameters(), fn, loader, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals[num_negative:], vals[num_negative:])
        assert torch.allclose(torch.ones(()), vals[:num_negative])

        assert torch.allclose(torch.zeros(()), vecs[:, :num_negative])
        _assert_eigvecs_allclose(ref_vecs[:, num_negative:], vecs[:, num_negative:])
//  [python]
def test_hessian_eigs_positive_lanczos(self):
        "Compare `eff_dim.hessian_eigs_positive_lanczos` with hessian then `torch.linalg.eigh`"
        model, fn, loader = self.example_model_loss_data()
        M = eff_dim.hessian(model.parameters(), fn, loader)
        ref_vals, ref_vecs = torch.linalg.eigh(M)
        num_negative = (ref_vals < 0).sum().item()

        # All values/vectors:
        vals, vecs = eff_dim.hessian_eigs_positive_lanczos(
            model.parameters(), fn, loader, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals[num_negative:], vals[num_negative:])
        assert torch.allclose(torch.ones(()), vals[:num_negative])

        assert torch.allclose(torch.zeros(()), vecs[:, :num_negative])
        _assert_eigvecs_allclose(ref_vecs[:, num_negative:], vecs[:, num_negative:])
//  [python]
def test_PD_hessian_eigs(self, D=4):
        "Like `test_hessian_eigs_pozitive_lanczos` but for a positive-definite Hessian"
        model, fn, loader = self.example_model_loss_data(linear_only=True)

        M = eff_dim.hessian(model.parameters(), fn, loader)
        ref_vals, ref_vecs = torch.symeig(M, eigenvectors=True)
        num_negative = (ref_vals < 0).sum().item()

        # All values/vectors:
        vals, vecs = eff_dim.hessian_eigs_positive_lanczos(
            model.parameters(), fn, loader, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals[num_negative:], vals[num_negative:], rtol=1e-4)
        assert torch.allclose(torch.ones(()), vals[:num_negative])
        _assert_eigvecs_allclose(ref_vecs[:, num_negative:], vecs[:, num_negative:], rtol=1e-3)
//  [python]
def test_PD_hessian_eigs(self, D=4):
        "Like `test_hessian_eigs_pozitive_lanczos` but for a positive-definite Hessian"
        model, fn, loader = self.example_model_loss_data(linear_only=True)

        M = eff_dim.hessian(model.parameters(), fn, loader)
        ref_vals, ref_vecs = torch.linalg.eigh(M)
        num_negative = (ref_vals < 0).sum().item()

        # All values/vectors:
        vals, vecs = eff_dim.hessian_eigs_positive_lanczos(
            model.parameters(), fn, loader, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals[num_negative:], vals[num_negative:], rtol=1e-4)
        assert torch.allclose(torch.ones(()), vals[:num_negative])
        _assert_eigvecs_allclose(ref_vecs[:, num_negative:], vecs[:, num_negative:], rtol=1e-3)
//  [python]
def test_forward():
    torch.manual_seed(42)
    N = 100
    tol = 1E-8
    dtype = torch.float64
    A = torch.randn(N, N, dtype=dtype)
    A = A+A.t()

    w, v = torch.symeig(A, eigenvectors=True)
    idx = torch.argmax(w.abs())

    v_exact = v[:, idx] 
    v_exact = v_exact[0].sign() * v_exact
    
    x0 = torch.rand(N, dtype=dtype)
    x0 = x0/x0.norm()
    x = FixedPoint.apply(A, x0, tol)

    assert(torch.allclose(v_exact, x, rtol=tol, atol=tol))
//  [python]
def test_forward():
    torch.manual_seed(42)
    N = 100
    tol = 1E-8
    dtype = torch.float64
    A = torch.randn(N, N, dtype=dtype)
    A = A+A.t()

    w, v = torch.linalg.eigh(A)
    idx = torch.argmax(w.abs())

    v_exact = v[:, idx] 
    v_exact = v_exact[0].sign() * v_exact
    
    x0 = torch.rand(N, dtype=dtype)
    x0 = x0/x0.norm()
    x = FixedPoint.apply(A, x0, tol)

    assert(torch.allclose(v_exact, x, rtol=tol, atol=tol))
//  [python]
def forward(self, X):
        X_mean = X.mean(dim=0, keepdim=True)
        X_c = X - X_mean
        eigenvalues, eigenvectors = torch.symeig(X_c.t().matmul(X_c).cpu(), eigenvectors=True, upper=True)
        S, U = eigenvalues.to(X.device), eigenvectors.to(X.device)
        self.eig = eigenvalues.min()
        U_t = U.t()
        W = U.matmul(S.rsqrt().diag()).matmul(U_t)
        # with torch.no_grad:
        W_inv = U.detach().matmul(S.detach().sqrt().diag()).matmul(U_t.detach())
        # breakpoint()
        X = ( X_c.matmul(W).matmul(W_inv) + X_mean.detach() - self.running_mean ).matmul(self.running_W)
        self.running_mean = self.momentum * X_mean.detach() + (1 - self.momentum) * self.running_mean
        self.running_W = self.momentum * W.detach() + (1 - self.momentum) * self.running_W
        return X
//  [python]
def forward(self, X):
        X_mean = X.mean(dim=0, keepdim=True)
        X_c = X - X_mean
        eigenvalues, eigenvectors = torch.linalg.eigh(X_c.t().matmul(X_c).cpu(), UPLO='U')
        S, U = eigenvalues.to(X.device), eigenvectors.to(X.device)
        self.eig = eigenvalues.min()
        U_t = U.t()
        W = U.matmul(S.rsqrt().diag()).matmul(U_t)
        # with torch.no_grad:
        W_inv = U.detach().matmul(S.detach().sqrt().diag()).matmul(U_t.detach())
        # breakpoint()
        X = ( X_c.matmul(W).matmul(W_inv) + X_mean.detach() - self.running_mean ).matmul(self.running_W)
        self.running_mean = self.momentum * X_mean.detach() + (1 - self.momentum) * self.running_mean
        self.running_W = self.momentum * W.detach() + (1 - self.momentum) * self.running_W
        return X
//  [python]
def __init__(
        self,
        mean: Tensor,
        cov: Tensor,
        seed: Optional[int] = None,
        inv_transform: bool = False,
    ) -> None:
        r"""Engine for qMC sampling from a multivariate Normal `N(\mu, \Sigma)`.

        Args:
            mean: The mean vector.
            cov: The covariance matrix.
            seed: The seed with which to seed the random number generator of the
                underlying SobolEngine.
            inv_transform: If True, use inverse transform instead of Box-Muller.
        """
        # validate inputs
        if not cov.shape[0] == cov.shape[1]:
            raise ValueError("Covariance matrix is not square.")
        if not mean.shape[0] == cov.shape[0]:
            raise ValueError("Dimension mismatch between mean and covariance.")
        if not torch.allclose(cov, cov.transpose(-1, -2)):
            raise ValueError("Covariance matrix is not symmetric.")
        self._mean = mean
        self._normal_engine = NormalQMCEngine(
            d=mean.shape[0], seed=seed, inv_transform=inv_transform
        )
        # compute Cholesky decomp; if it fails, do the eigendecomposition
        try:
            self._corr_matrix = torch.cholesky(cov).transpose(-1, -2)
        except RuntimeError:
            eigval, eigvec = torch.symeig(cov, eigenvectors=True)
            if not torch.all(eigval >= -1e-8):
                raise ValueError("Covariance matrix not PSD.")
            eigval_root = eigval.clamp_min(0.0).sqrt()
            self._corr_matrix = (eigvec * eigval_root).transpose(-1, -2)
//  [python]
def __init__(
        self,
        mean: Tensor,
        cov: Tensor,
        seed: Optional[int] = None,
        inv_transform: bool = False,
    ) -> None:
        r"""Engine for qMC sampling from a multivariate Normal `N(\mu, \Sigma)`.

        Args:
            mean: The mean vector.
            cov: The covariance matrix.
            seed: The seed with which to seed the random number generator of the
                underlying SobolEngine.
            inv_transform: If True, use inverse transform instead of Box-Muller.
        """
        # validate inputs
        if not cov.shape[0] == cov.shape[1]:
            raise ValueError("Covariance matrix is not square.")
        if not mean.shape[0] == cov.shape[0]:
            raise ValueError("Dimension mismatch between mean and covariance.")
        if not torch.allclose(cov, cov.transpose(-1, -2)):
            raise ValueError("Covariance matrix is not symmetric.")
        self._mean = mean
        self._normal_engine = NormalQMCEngine(
            d=mean.shape[0], seed=seed, inv_transform=inv_transform
        )
        # compute Cholesky decomp; if it fails, do the eigendecomposition
        try:
            self._corr_matrix = torch.cholesky(cov).transpose(-1, -2)
        except RuntimeError:
            eigval, eigvec = torch.linalg.eigh(cov)
            if not torch.all(eigval >= -1e-8):
                raise ValueError("Covariance matrix not PSD.")
            eigval_root = eigval.clamp_min(0.0).sqrt()
            self._corr_matrix = (eigvec * eigval_root).transpose(-1, -2)
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b,d,d)
    for batch_idx in range(X.shape[0]):
        e,v = torch.symeig(X[batch_idx,:,:].squeeze(), True)
        bv[batch_idx,:,:] = v
    bv = bv.cuda()
    return bv
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b,d,d)
    for batch_idx in range(X.shape[0]):
        e,v = torch.linalg.eigh(X[batch_idx,:,:].squeeze())
        bv[batch_idx,:,:] = v
    bv = bv.cuda()
    return bv
//  [python]
def symeig(A, largest=False, eigenvectors=True):
    # type: (Tensor, Optional[bool], Optional[bool]) -> Tuple[Tensor, Tensor]
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    if eigenvectors is None:
        eigenvectors = True
    E, Z = torch.symeig(A, eigenvectors, True)
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def symeig(A, largest=False, eigenvectors=True):
    # type: (Tensor, Optional[bool], Optional[bool]) -> Tuple[Tensor, Tensor]
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    if eigenvectors is None:
        eigenvectors = True
    E, Z = torch.linalg.eigh(A, eigenvectors=eigenvectors)
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def logm_torch(x):
    """
    This function computes the logarithm of a matrix.

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

    Returns
    -------
    :return: logm(x)
    """
    eigendecomposition = torch.symeig(x, eigenvectors=True)

    eigenvectors = eigendecomposition.eigenvectors
    log_eigenvalues = torch.log(eigendecomposition.eigenvalues)  # Assume real eigenvalues (first column only)

    return torch.mm(eigenvectors, torch.mm(torch.diag(log_eigenvalues), torch.inverse(eigenvectors)))
//  [python]
def logm_torch(x):
    """
    This function computes the logarithm of a matrix.

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

    Returns
    -------
    :return: logm(x)
    """
    eigendecomposition = torch.linalg.eigh(x)

    eigenvectors = eigendecomposition.eigenvectors
    log_eigenvalues = torch.log(eigendecomposition.eigenvalues)  # Assume real eigenvalues (first column only)

    return torch.mm(eigenvectors, torch.mm(torch.diag(log_eigenvalues), torch.inverse(eigenvectors)))
//  [python]
def get_singular_values(cov, eps=1e-5):
    #input: covariance matrix C*C
    #output: singular values in increasing order: C
    C,_ = cov.shape
    cov += eps*torch.eye(C).cuda() #for numerical stability
    s, _ = torch.symeig(cov)
    return s
//  [python]
def get_singular_values(cov, eps=1e-5):
    #input: covariance matrix C*C
    #output: singular values in increasing order: C
    C,_ = cov.shape
    cov += eps*torch.eye(C).cuda() #for numerical stability
    s, _ = torch.linalg.eigh(cov)
    return s
//  [python]
def sqrtm_torch(x):
    """
    This function computes the square root of a matrix.

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

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

    eigenvectors = eigendecomposition.eigenvectors
    sqrt_eigenvalues = torch.sqrt(eigendecomposition.eigenvalues)  # Assume real eigenvalues (first column only)

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

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

    Returns
    -------
    :return: sqrtm(x)
    """
    eigendecomposition = torch.linalg.eigh(x)

    eigenvectors = eigendecomposition.eigenvectors
    sqrt_eigenvalues = torch.sqrt(eigendecomposition.eigenvalues)  # Assume real eigenvalues (first column only)

    return torch.mm(eigenvectors, torch.mm(torch.diag(sqrt_eigenvalues), torch.inverse(eigenvectors)))
//  [python]
def train(self, x, inplace=False):
    """
      train PCA with 'x'
      x: torch.Tensor, shape : [d_vec, n_sample]
      inplace: bool, default : False
        if True, reduce the memory consumption with the cost of extra computation
    """
    assert x.shape[0] >= self.n_components

    mean = x.mean(dim=1, keepdim=True) #[d_vec, 1]
    if inplace:
      x.sub_(mean)
    else:
      x = x - mean
    x_cov = self.covar(x, rowvar=True, meaned=True)
    if inplace:
      x.add_(mean)
    eig_val, eig_vec = torch.symeig(x_cov, eigenvectors=True, upper=False)
    sorted_eig_val, sorted_index = eig_val.sort(descending=True)
    sorted_eig_vec = eig_vec[:, sorted_index]
    components = sorted_eig_vec[:, :self.n_components].T
    self.register_buffer("_components", components)
    self.register_buffer("_mean", mean)
//  [python]
def train(self, x, inplace=False):
    """
      train PCA with 'x'
      x: torch.Tensor, shape : [d_vec, n_sample]
      inplace: bool, default : False
        if True, reduce the memory consumption with the cost of extra computation
    """
    assert x.shape[0] >= self.n_components

    mean = x.mean(dim=1, keepdim=True) #[d_vec, 1]
    if inplace:
      x.sub_(mean)
    else:
      x = x - mean
    x_cov = self.covar(x, rowvar=True, meaned=True)
    if inplace:
      x.add_(mean)
    eig_val, eig_vec = torch.linalg.eigh(x_cov)
    sorted_eig_val, sorted_index = eig_val.sort(descending=True)
    sorted_eig_vec = eig_vec[:, sorted_index]
    components = sorted_eig_vec[:, :self.n_components].T
    self.register_buffer("_components", components)
    self.register_buffer("_mean", mean)
//  [python]
def test_categorical_hessian(fcat):
    # cannot test, algebraic implementation.
    # make sure it runs and is decomposable!
    lh = CategoricalLh()
    H = lh.Hessian(fcat)
    ls = torch.symeig(H)[0]
    assert ls.shape[0] == len(fcat) and ls.shape[1] == len(fcat[0])
    assert H.shape[1] == H.shape[2]
//  [python]
def test_categorical_hessian(fcat):
    # cannot test, algebraic implementation.
    # make sure it runs and is decomposable!
    lh = CategoricalLh()
    H = lh.Hessian(fcat)
    ls = torch.linalg.eigh(H)[0]
    assert ls.shape[0] == len(fcat) and ls.shape[1] == len(fcat[0])
    assert H.shape[1] == H.shape[2]
//  [python]
def spd_to_so_and_euclidean_torch(x):
    """
    Compute the points on the SPD manifold as a product of SO and euclidean manifolds.

    Parameters
    ----------
    :param x: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)

    Returns
    -------
    :return: (set of Euclidean points, set of SO points) corresponding to the SPD matrices
    """
    # Decompose SPD matrices in SO(n) and Euclidean parts using eigenvalue decompositions
    eigendecomposition1 = torch.symeig(x, eigenvectors=True)
    eigenvecs = eigendecomposition1.eigenvectors
    det_xso = torch.det(eigenvecs)  # Only matrices with det=1 belong to SO, so we swap the sign if det=-1
    # x_so[det_xso < 0.] *= -1
    x_so = det_xso.unsqueeze(-1).unsqueeze(-1) * eigenvecs
    x_euclidean = eigendecomposition1.eigenvalues  # Assume real eigenvalues (first column only)

    return x_euclidean, x_so
//  [python]
def spd_to_so_and_euclidean_torch(x):
    """
    Compute the points on the SPD manifold as a product of SO and euclidean manifolds.

    Parameters
    ----------
    :param x: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)

    Returns
    -------
    :return: (set of Euclidean points, set of SO points) corresponding to the SPD matrices
    """
    # Decompose SPD matrices in SO(n) and Euclidean parts using eigenvalue decompositions
    eigendecomposition1 = torch.linalg.eigh(x)
    eigenvecs = eigendecomposition1.eigenvectors
    det_xso = torch.det(eigenvecs)  # Only matrices with det=1 belong to SO, so we swap the sign if det=-1
    # x_so[det_xso < 0.] *= -1
    x_so = det_xso.unsqueeze(-1).unsqueeze(-1) * eigenvecs
    x_euclidean = eigendecomposition1.eigenvalues  # Assume real eigenvalues (first column only)

    return x_euclidean, x_so
//  [python]
def torch_func(b):
    c, d, h, w = 9, 10, 10, 10
    return timeit.timeit(lambda: torch.symeig(rand_torch_input((b, c, d, h, w)), eigenvectors=True), number=5) / 5
//  [python]
def torch_func(b):
    c, d, h, w = 9, 10, 10, 10
    return timeit.timeit(lambda: torch.linalg.eigh(rand_torch_input((b, c, d, h, w))), number=5) / 5
//  [python]
def create_frame(self, X, mask):
        mask = mask.unsqueeze(-1)
        center = (X * mask).sum(dim=1) / mask.sum(dim=1)
        X = X - center.unsqueeze(1) * mask  # [B,N,3]
        C = torch.bmm(X.transpose(1,2), X)  # [B,3,3] (Cov)
        _, V = torch.symeig(C.detach(), True)  # [B,3,3]
        F_ops = self.ops.unsqueeze(1).unsqueeze(0) * V.unsqueeze(1)  # [1,8,1,3] x [B,1,3,3] -> [B,8,3,3]
        h = torch.einsum('boij,bpj->bopi', F_ops.transpose(2,3), X)  # transpose is inverse [B,8,N,3]
        h = h.view(X.size(0) * 8, X.size(1), 3)
        return h, F_ops.detach(), center
//  [python]
def create_frame(self, X, mask):
        mask = mask.unsqueeze(-1)
        center = (X * mask).sum(dim=1) / mask.sum(dim=1)
        X = X - center.unsqueeze(1) * mask  # [B,N,3]
        C = torch.bmm(X.transpose(1,2), X)  # [B,3,3] (Cov)
        _, V = torch.linalg.eigh(C.detach())  # [B,3,3]
        F_ops = self.ops.unsqueeze(1).unsqueeze(0) * V.unsqueeze(1)  # [1,8,1,3] x [B,1,3,3] -> [B,8,3,3]
        h = torch.einsum('boij,bpj->bopi', F_ops.transpose(2,3), X)  # transpose is inverse [B,8,N,3]
        h = h.view(X.size(0) * 8, X.size(1), 3)
        return h, F_ops.detach(), center
//  [python]
def torch_func_gpu(b):
    c, d, h, w = 9, 10, 10, 10
    return timeit.timeit(lambda: torch.symeig(rand_torch_input((b, c, d, h, w)).cuda(), eigenvectors=True),
                         number=5) / 5
//  [python]
def torch_func_gpu(b):
    c, d, h, w = 9, 10, 10, 10
    return timeit.timeit(lambda: torch.linalg.eigh(rand_torch_input((b, c, d, h, w)).cuda()),
                         number=5) / 5
//  [python]
def loss_pos_matrix_eigen_values(A):
    # A: bs x d x d
    eigv = torch.symeig(A, eigenvectors=True)[0].view(-1)
    negative_index = eigv.detach().cpu().numpy() < 0
    negative_eigv = eigv[negative_index]
    return negative_eigv.norm()
//  [python]
def loss_pos_matrix_eigen_values(A):
    # A: bs x d x d
    eigv = torch.linalg.eigh(A)[0].view(-1)
    negative_index = eigv.detach().cpu().numpy() < 0
    negative_eigv = eigv[negative_index]
    return negative_eigv.norm()
//  [python]
def fit(self, data, weight=None):

        assert data.ndim == 2 and data.shape[1] == self.ndim_data
        if weight is not None:
            raise NotImplementedError

        with torch.no_grad():
            self.mean[:] = torch.mean(data, dim=0)
            data0 = data - self.mean
            covariance = data0.T @ data0 / (data0.shape[0]-1)
            D, E = torch.symeig(covariance, eigenvectors=True)
            self.D[:] = torch.flip(D, dims=(0,))
            self.E[:] = torch.flip(E, dims=(1,))

            return self
//  [python]
def fit(self, data, weight=None):

        assert data.ndim == 2 and data.shape[1] == self.ndim_data
        if weight is not None:
            raise NotImplementedError

        with torch.no_grad():
            self.mean[:] = torch.mean(data, dim=0)
            data0 = data - self.mean
            covariance = data0.T @ data0 / (data0.shape[0]-1)
            D, E = torch.linalg.eigh(covariance)
            self.D[:] = torch.flip(D, dims=(0,))
            self.E[:] = torch.flip(E, dims=(1,))

            return self
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    # print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  # (B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  # (B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  # (B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask > 1] = 1

        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.symeig(L, eigenvectors=True)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    # print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  # (B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  # (B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  # (B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask > 1] = 1

        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.linalg.eigh(L)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def chiF_perturbation(model):
    """
        Compute chi_F using the full spectrum perturbation formula:
    chi_F = \sum_{n \neq 0} \frac{ |<psi_n(g)| \partial H / \partial g |psi_0(g)>|^2 }
                                 {(E_0(g) - E_n(g))^2}
    """
    Es, psis = torch.symeig(model.Hmatrix, eigenvectors=True)
    psi0 = psis[:, 0]

    model.setpHpg()
    numerators = psi0.matmul(model.pHpgmatrix).matmul(psis)[1:] ** 2
    denominators = (Es[0] - Es[1:]) ** 2
    chiF = (numerators / denominators).sum().item()
    return chiF
//  [python]
def chiF_perturbation(model):
    """
        Compute chi_F using the full spectrum perturbation formula:
    chi_F = \sum_{n \neq 0} \frac{ |<psi_n(g)| \partial H / \partial g |psi_0(g)>|^2 }
                                 {(E_0(g) - E_n(g))^2}
    """
    Es, psis = torch.linalg.eigh(model.Hmatrix)
    psi0 = psis[:, 0]

    model.setpHpg()
    numerators = psi0.matmul(model.pHpgmatrix).matmul(psis)[1:] ** 2
    denominators = (Es[0] - Es[1:]) ** 2
    chiF = (numerators / denominators).sum().item()
    return chiF
//  [python]
def compute_minimum_ratio(model, criterion, train, target):
    model.zero_grad()
    # compute Hessian matrix
    # save the gradient of each layer
    with autograd_lib.module_hook(save_activations):
        output = model(train)
        loss = criterion(output, target)

    # compute Hessian according to the gradient value stored in the previous step
    with autograd_lib.module_hook(compute_hess):
        autograd_lib.backward_hessian(output, loss='LeastSquares')

    # [0].shape: [1, 128, 1, 128] & [1].shape: [128, 1, 128, 1]
    layer_hess = list(hess.values())
    minimum_ratio = []

    # compute eigenvalues of the Hessian matrix
    for h in layer_hess:
        size = h.shape[0] * h.shape[1]
        h = h.reshape(size, size)
        # torch.symeig() returns eigenvalues and eigenvectors of a real symmetric matrix
        h_eig = torch.symeig(h).eigenvalues
        num_greater = torch.sum(h_eig > 0).item()
        minimum_ratio.append(num_greater / len(h_eig))

    ratio_mean = np.mean(minimum_ratio)  # compute mean of minimum ratio

    return ratio_mean
//  [python]
def compute_minimum_ratio(model, criterion, train, target):
    model.zero_grad()
    # compute Hessian matrix
    # save the gradient of each layer
    with autograd_lib.module_hook(save_activations):
        output = model(train)
        loss = criterion(output, target)

    # compute Hessian according to the gradient value stored in the previous step
    with autograd_lib.module_hook(compute_hess):
        autograd_lib.backward_hessian(output, loss='LeastSquares')

    # [0].shape: [1, 128, 1, 128] & [1].shape: [128, 1, 128, 1]
    layer_hess = list(hess.values())
    minimum_ratio = []

    # compute eigenvalues of the Hessian matrix
    for h in layer_hess:
        size = h.shape[0] * h.shape[1]
        h = h.reshape(size, size)
        # torch.linalg.eigh returns eigenvalues and eigenvectors of a real symmetric matrix
        h_eig = torch.linalg.eigh(h).eigenvalues
        num_greater = torch.sum(h_eig > 0).item()
        minimum_ratio.append(num_greater / len(h_eig))

    ratio_mean = np.mean(minimum_ratio)  # compute mean of minimum ratio

    return ratio_mean
//  [python]
def _solve_from_interaction(self, interactions, T, mag_field=None):
        h_matrix, SpinOps = self.get_hamiltonian_matrix(interactions=interactions, mag_field=mag_field)
        sim_therm_data = ThObs(T=T)

        [ESpec, U] = torch.symeig(h_matrix, eigenvectors=True)
        Sztot = sum(SpinOps.SzP).double()
        Sxtot = sum(SpinOps.SxP).double()
        Mnn = torch.diag(U.transpose(1, 0).mm(Sztot).mm(U))
        Mnn_x = torch.diag(U.transpose(1, 0).mm(Sxtot).mm(U))
        MSqnn = torch.diag(U.transpose(1, 0).mm(Sztot).mm(Sztot).mm(U))
        for i in range(sim_therm_data.N):
            sim_therm_data.Z[i] = (torch.exp(-sim_therm_data.beta[i] * ESpec)).sum()

        for i in range(sim_therm_data.N):
            sim_therm_data.Fe[i, 0] = sim_therm_data.beta[i]
            sim_therm_data.Fe[i, 1] = -1 / sim_therm_data.beta[i] * \
                                              torch.log((torch.exp(-sim_therm_data.beta[i] * ESpec)).sum())
            sim_therm_data.M[i] = (Mnn * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
            sim_therm_data.M_paral[i] = (Mnn_x * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
            sim_therm_data.E[i] = (ESpec * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]

        for i in range(sim_therm_data.N):
            sim_therm_data.C[i] = ((ESpec ** 2 * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
                                   - sim_therm_data.E[i] ** 2) * sim_therm_data.beta[i] ** 2
            sim_therm_data.Chiz[i] = ((MSqnn * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
                                     - sim_therm_data.M[i] ** 2) * sim_therm_data.beta[i]

        sim_therm_data.C /= self.size
        sim_therm_data.Chiz /= self.size
        sim_therm_data.M /= self.size
        sim_therm_data.M_paral /= self.size

        return sim_therm_data
//  [python]
def _solve_from_interaction(self, interactions, T, mag_field=None):
        h_matrix, SpinOps = self.get_hamiltonian_matrix(interactions=interactions, mag_field=mag_field)
        sim_therm_data = ThObs(T=T)

        [ESpec, U] = torch.linalg.eigh(h_matrix)
        Sztot = sum(SpinOps.SzP).double()
        Sxtot = sum(SpinOps.SxP).double()
        Mnn = torch.diag(U.transpose(1, 0).mm(Sztot).mm(U))
        Mnn_x = torch.diag(U.transpose(1, 0).mm(Sxtot).mm(U))
        MSqnn = torch.diag(U.transpose(1, 0).mm(Sztot).mm(Sztot).mm(U))
        for i in range(sim_therm_data.N):
            sim_therm_data.Z[i] = (torch.exp(-sim_therm_data.beta[i] * ESpec)).sum()

        for i in range(sim_therm_data.N):
            sim_therm_data.Fe[i, 0] = sim_therm_data.beta[i]
            sim_therm_data.Fe[i, 1] = -1 / sim_therm_data.beta[i] * \
                                              torch.log((torch.exp(-sim_therm_data.beta[i] * ESpec)).sum())
            sim_therm_data.M[i] = (Mnn * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
            sim_therm_data.M_paral[i] = (Mnn_x * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
            sim_therm_data.E[i] = (ESpec * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]

        for i in range(sim_therm_data.N):
            sim_therm_data.C[i] = ((ESpec ** 2 * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
                                   - sim_therm_data.E[i] ** 2) * sim_therm_data.beta[i] ** 2
            sim_therm_data.Chiz[i] = ((MSqnn * torch.exp(-sim_therm_data.beta[i] * ESpec)).sum() / sim_therm_data.Z[i]
                                     - sim_therm_data.M[i] ** 2) * sim_therm_data.beta[i]

        sim_therm_data.C /= self.size
        sim_therm_data.Chiz /= self.size
        sim_therm_data.M /= self.size
        sim_therm_data.M_paral /= self.size

        return sim_therm_data
//  [python]
def forward(self,A,X):
    lr = A
    lr_dim = lr.shape[0]
    f = X
    eig_val_lr, U_lr = torch.symeig(lr, eigenvectors=True,upper=True)
    # U_lr = torch.abs(U_lr)
    eye_mat = torch.eye(lr_dim).type(torch.FloatTensor)
    s_d = torch.cat((eye_mat,eye_mat),0)
    
    a = torch.matmul(self.weights,s_d )
    b = torch.matmul(a ,torch.t(U_lr))
    f_d = torch.matmul(b ,f)
    f_d = torch.abs(f_d)
    self.f_d = f_d.fill_diagonal_(1)
    adj = normalize_adj_torch(self.f_d)
    X = torch.mm(adj, adj.t())
    X = (X + X.t())/2
    idx = torch.eye(320, dtype=bool)
    X[idx]=1
    return adj, torch.abs(X)
//  [python]
def forward(self,A,X):
    lr = A
    lr_dim = lr.shape[0]
    f = X
    eig_val_lr, U_lr = torch.linalg.eigh(lr)
    # U_lr = torch.abs(U_lr)
    eye_mat = torch.eye(lr_dim).type(torch.FloatTensor)
    s_d = torch.cat((eye_mat,eye_mat),0)
    
    a = torch.matmul(self.weights,s_d )
    b = torch.matmul(a ,torch.t(U_lr))
    f_d = torch.matmul(b ,f)
    f_d = torch.abs(f_d)
    self.f_d = f_d.fill_diagonal_(1)
    adj = normalize_adj_torch(self.f_d)
    X = torch.mm(adj, adj.t())
    X = (X + X.t())/2
    idx = torch.eye(320, dtype=bool)
    X[idx]=1
    return adj, torch.abs(X)
//  [python]
def __init__(self, lamda, random_seed_pair=(None, None)):
        self.lamda = lamda
        self.n_vertices = np.array([2] * CONTAMINATION_N_STAGES)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1])], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
        # In all evaluation, the same sampled values are used.
        self.init_Z, self.lambdas, self.gammas = generate_contamination_dynamics(random_seed_pair[0])
//  [python]
def __init__(self, lamda, random_seed_pair=(None, None)):
        self.lamda = lamda
        self.n_vertices = np.array([2] * CONTAMINATION_N_STAGES)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1])], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.linalg.eigh(laplacian)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
        # In all evaluation, the same sampled values are used.
        self.init_Z, self.lambdas, self.gammas = generate_contamination_dynamics(random_seed_pair[0])
//  [python]
def __init__(self, lamda, random_seed_pair=(None, None)):
        self.lamda = lamda
        self.n_vertices = np.array([2] * ISING_N_EDGES)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
        interaction = generate_ising_interaction(ISING_GRID_H, ISING_GRID_W, random_seed_pair[0])
        self.interaction = interaction[0].numpy(), interaction[1].numpy()
        self.covariance, self.partition_original = spin_covariance(self.interaction, (ISING_GRID_H, ISING_GRID_W))
//  [python]
def __init__(self, lamda, random_seed_pair=(None, None)):
        self.lamda = lamda
        self.n_vertices = np.array([2] * ISING_N_EDGES)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.linalg.eigh(laplacian)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
        interaction = generate_ising_interaction(ISING_GRID_H, ISING_GRID_W, random_seed_pair[0])
        self.interaction = interaction[0].numpy(), interaction[1].numpy()
        self.covariance, self.partition_original = spin_covariance(self.interaction, (ISING_GRID_H, ISING_GRID_W))
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.
        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.
        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.symeig(matrix, eigenvectors=True)
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.
        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.
        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.linalg.eigh(matrix)
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def hosvd(x):
    """
    :param x: 待分解的张量
    :return G: 核张量
    :return U: 变换矩阵
    :return lm: 各个键约化矩阵的本征谱
    """
    ndim = x.ndim
    U = list()  # 用于存取各个变换矩阵
    lm = list()  # 用于存取各个键约化矩阵的本征谱
    x = tc.from_numpy(x)
    for n in range(ndim):
        index = list(range(ndim))
        index.pop(n)
        _mat = tc.tensordot(x, x, [index, index])
        _lm, _U = tc.symeig(_mat, eigenvectors=True)
        lm.append(_lm.numpy())
        U.append(_U)
    # 计算核张量
    G = tucker_product(x, U)
    U1 = [u.numpy() for u in U]
    return G, U1, lm
//  [python]
def hosvd(x):
    """
    :param x: 待分解的张量
    :return G: 核张量
    :return U: 变换矩阵
    :return lm: 各个键约化矩阵的本征谱
    """
    ndim = x.ndim
    U = list()  # 用于存取各个变换矩阵
    lm = list()  # 用于存取各个键约化矩阵的本征谱
    x = tc.from_numpy(x)
    for n in range(ndim):
        index = list(range(ndim))
        index.pop(n)
        _mat = tc.tensordot(x, x, [index, index])
        _lm, _U = tc.linalg.eigh(_mat)
        lm.append(_lm.numpy())
        U.append(_U)
    # 计算核张量
    G = tucker_product(x, U)
    U1 = [u.numpy() for u in U]
    return G, U1, lm
//  [python]
def estimate_perpendicular(pc, k, sigma=0.01, clip=0.05):
    with torch.no_grad():
        # pc : [b, 3, n]
        b,_,n=pc.size()
        inter_KNN = knn_points(pc.permute(0,2,1), pc.permute(0,2,1), K=k+1) #[dists:[b,n,k+1], idx:[b,n,k+1]]
        nn_pts = knn_gather(pc.permute(0,2,1), inter_KNN.idx).permute(0,3,1,2)[:,:,:,1:].contiguous() # [b, 3, n ,k]

        # get covariance matrix and smallest eig-vector of individual point
        perpendi_vector_1 = []
        perpendi_vector_2 = []
        for i in range(b):
            curr_point_set = nn_pts[i].detach().permute(1,0,2) #curr_point_set:[n, 3, k]
            curr_point_set_mean = torch.mean(curr_point_set, dim=2, keepdim=True) #curr_point_set_mean:[n, 3, 1]
            curr_point_set = curr_point_set - curr_point_set_mean #curr_point_set:[n, 3, k]
            curr_point_set_t = curr_point_set.permute(0,2,1) #curr_point_set_t:[n, k, 3]
            fact = 1.0 / (k-1)
            cov_mat = fact * torch.bmm(curr_point_set, curr_point_set_t) #curr_point_set_t:[n, 3, 3]
            eigenvalue, eigenvector=torch.symeig(cov_mat, eigenvectors=True)    # eigenvalue:[n, 3], eigenvector:[n, 3, 3]

            larger_dim_idx = torch.topk(eigenvalue, 2, dim=1, largest=True, sorted=False, out=None)[1] # eigenvalue:[n, 2]

            persample_perpendi_vector_1 = torch.gather(eigenvector, 2, larger_dim_idx[:,0].unsqueeze(1).unsqueeze(2).expand(n, 3, 1)).squeeze() #persample_perpendi_vector_1:[n, 3]
            persample_perpendi_vector_2 = torch.gather(eigenvector, 2, larger_dim_idx[:,1].unsqueeze(1).unsqueeze(2).expand(n, 3, 1)).squeeze() #persample_perpendi_vector_2:[n, 3]

            perpendi_vector_1.append(persample_perpendi_vector_1.permute(1,0))
            perpendi_vector_2.append(persample_perpendi_vector_2.permute(1,0))

        perpendi_vector_1 = torch.stack(perpendi_vector_1, 0) #perpendi_vector_1:[b, 3, n]
        perpendi_vector_2 = torch.stack(perpendi_vector_2, 0) #perpendi_vector_1:[b, 3, n]

        aux_vector1 = sigma * torch.randn(b,n).unsqueeze(1).cuda() #aux_vector1:[b, 1, n]
        aux_vector2 = sigma * torch.randn(b,n).unsqueeze(1).cuda() #aux_vector2:[b, 1, n]

    return torch.clamp(perpendi_vector_1*aux_vector1, -1*clip, clip) + torch.clamp(perpendi_vector_2*aux_vector2, -1*clip, clip)
//  [python]
def estimate_perpendicular(pc, k, sigma=0.01, clip=0.05):
    with torch.no_grad():
        # pc : [b, 3, n]
        b,_,n=pc.size()
        inter_KNN = knn_points(pc.permute(0,2,1), pc.permute(0,2,1), K=k+1) #[dists:[b,n,k+1], idx:[b,n,k+1]]
        nn_pts = knn_gather(pc.permute(0,2,1), inter_KNN.idx).permute(0,3,1,2)[:,:,:,1:].contiguous() # [b, 3, n ,k]

        # get covariance matrix and smallest eig-vector of individual point
        perpendi_vector_1 = []
        perpendi_vector_2 = []
        for i in range(b):
            curr_point_set = nn_pts[i].detach().permute(1,0,2) #curr_point_set:[n, 3, k]
            curr_point_set_mean = torch.mean(curr_point_set, dim=2, keepdim=True) #curr_point_set_mean:[n, 3, 1]
            curr_point_set = curr_point_set - curr_point_set_mean #curr_point_set:[n, 3, k]
            curr_point_set_t = curr_point_set.permute(0,2,1) #curr_point_set_t:[n, k, 3]
            fact = 1.0 / (k-1)
            cov_mat = fact * torch.bmm(curr_point_set, curr_point_set_t) #curr_point_set_t:[n, 3, 3]
            eigenvalue, eigenvector=torch.linalg.eigh(cov_mat)    # eigenvalue:[n, 3], eigenvector:[n, 3, 3]

            larger_dim_idx = torch.topk(eigenvalue, 2, dim=1, largest=True, sorted=False, out=None)[1] # eigenvalue:[n, 2]

            persample_perpendi_vector_1 = torch.gather(eigenvector, 2, larger_dim_idx[:,0].unsqueeze(1).unsqueeze(2).expand(n, 3, 1)).squeeze() #persample_perpendi_vector_1:[n, 3]
            persample_perpendi_vector_2 = torch.gather(eigenvector, 2, larger_dim_idx[:,1].unsqueeze(1).unsqueeze(2).expand(n, 3, 1)).squeeze() #persample_perpendi_vector_2:[n, 3]

            perpendi_vector_1.append(persample_perpendi_vector_1.permute(1,0))
            perpendi_vector_2.append(persample_perpendi_vector_2.permute(1,0))

        perpendi_vector_1 = torch.stack(perpendi_vector_1, 0) #perpendi_vector_1:[b, 3, n]
        perpendi_vector_2 = torch.stack(perpendi_vector_2, 0) #perpendi_vector_1:[b, 3, n]

        aux_vector1 = sigma * torch.randn(b,n).unsqueeze(1).cuda() #aux_vector1:[b, 1, n]
        aux_vector2 = sigma * torch.randn(b,n).unsqueeze(1).cuda() #aux_vector2:[b, 1, n]

    return torch.clamp(perpendi_vector_1*aux_vector1, -1*clip, clip) + torch.clamp(perpendi_vector_2*aux_vector2, -1*clip, clip)
//  [python]
def augment_on_the_fly(self, g: Graph) -> Graph:
        x, edge_index, edge_prob = g.unfold()
        x = x.to(self.device)
        ori_adj = get_adj_tensor(edge_index.cpu()).to(self.device)
        # ori_adj = to_dense_adj(edge_index)
        
        nnodes = ori_adj.shape[0]
                
        adj_changes = Parameter(torch.FloatTensor(int(nnodes*(nnodes-1)/2)), requires_grad=True).to(self.device)
        torch.nn.init.uniform_(adj_changes, 0.0, 0.001)
        
        ori_adj_norm = get_normalize_adj_tensor(ori_adj, device=self.device)
        # ori_e = torch.linalg.eigvalsh(ori_adj_norm)
        ori_e, ori_v = torch.symeig(ori_adj_norm, eigenvectors=True)
        eigen_norm = torch.norm(ori_e)
        
        # print(ori_adj.shape, ori_adj_norm.shape)
        # exit('')
        
        n_perturbations = int(self.ratio * (ori_adj.sum()/2))
        with tqdm(total=self.iteration, desc='Spectral Augment') as pbar:
            for t in range(1, self.iteration+1):
                modified_adj = self.get_modified_adj(ori_adj, self.reshape_m(nnodes, adj_changes))
                
                # add noise to make the graph asymmetric
                modified_adj_noise = modified_adj
                modified_adj_noise = self.add_random_noise(modified_adj)
                adj_norm_noise = get_normalize_adj_tensor(modified_adj_noise, device=self.device)
                # e = torch.linalg.eigvalsh(adj_norm_noise)
                e, v = torch.symeig(adj_norm_noise, eigenvectors=True)
                eigen_self = torch.norm(e)
                
                # spectral distance
                eigen_mse = torch.norm(ori_e-e)
                
                if self.dis_type == 'l2':
                    reg_loss = eigen_mse / eigen_norm
                elif self.dis_type == 'normDiv':
                    reg_loss = eigen_self / eigen_norm
                else:
                    exit(f'unknown distance metric: {self.dis_type}')

                self.loss = reg_loss
                
                adj_grad = torch.autograd.grad(self.loss, adj_changes)[0]

                lr = self.lr / np.sqrt(t+1)
                adj_changes.data.add_(lr * adj_grad)
                    
                before_p = torch.clamp(adj_changes, 0, 1).sum()
                self.projection(n_perturbations, adj_changes)
                after_p = torch.clamp(adj_changes, 0, 1).sum()
                
                pbar.set_postfix({'reg_loss': reg_loss.item(), 'eigen_mse': eigen_mse.item(), 'before_p': before_p.item(), 'after_p': after_p.item()})
                pbar.update()
        
        adj_changes = self.random_sample(adj_changes)
        
        modified_adj = self.get_modified_adj(ori_adj, self.reshape_m(nnodes, adj_changes)).detach()
        self.check_adj_tensor(modified_adj)
            
        edge_index, _ = dense_to_sparse(modified_adj)
        return Graph(x=x, edge_index=edge_index)