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

        self.n_eval += y.shape[0]

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

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

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

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

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

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

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

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

        self.C = (1 - self.c_r1 - self.c_mu) * self.C + self.c_r1 * C_r1 + self.c_mu * C_mu
//  [python]
def test_eig_with_nan(self, device, dtype):
        for val in [np.inf, np.nan]:
            for batch_dim in [(), (10,)]:
                a = make_tensor((*batch_dim, 5, 5), device=device, dtype=dtype)
                a[..., -1, -1] = val

                with self.assertRaisesRegex(RuntimeError, "torch.linalg.eig: input tensor should not"):
                    t.eig(a)
//  [python]
def test_eig_with_nan(self, device, dtype):
        for val in [np.inf, np.nan]:
            for batch_dim in [(), (10,)]:
                a = make_tensor((*batch_dim, 5, 5), device=device, dtype=dtype)
                a[..., -1, -1] = val

                with self.assertRaisesRegex(RuntimeError, "torch.linalg.eig: input tensor should not"):
                    torch.linalg.eig(a)
//  [python]
def check_pos_def(self, inputs):
        assert torch.all(inputs == inputs.mT), "Input matrix is not symmetric."
        assert torch.all(torch.symeig(inputs)[0] >= 0), (
            "Input matrix is not positive semi-definite in order to perform Cholesky decomposition"
        )
//  [python]
def check_pos_def(self, inputs):
        assert torch.all(inputs == inputs.mT), "Input matrix is not symmetric."
        assert torch.all(torch.linalg.eig(inputs)[0].real >= 0), (
            "Input matrix is not positive semi-definite in order to perform Cholesky decomposition"
        )
//  [python]
def _eig(self, cov: torch.Tensor, dtype: torch.dtype = torch.float32):
        """
        Computes eigenvalues and eigenvectors. Returns only the real part of the complex numbers.
        """
        eig_values, eig_vectors = torch.eig(cov)
        eig_values, eig_vectors = eig_values.unsqueeze(-1).type(dtype), eig_vectors.type(dtype)
        return eig_values, eig_vectors
//  [python]
def _eig(self, cov: torch.Tensor, dtype: torch.dtype = torch.float32):
        """
        Computes eigenvalues and eigenvectors. Returns only the real part of the complex numbers.
        """
        eig_values, eig_vectors = torch.linalg.eig(cov)
        eig_values, eig_vectors = eig_values.unsqueeze(-1).type(dtype), eig_vectors.type(dtype)
        return eig_values, eig_vectors
//  [python]
def kw_sq_to_ks(kw_sq, kernel_type, is_scalar=False):
    if kernel_type[-2:] == '_a' and not is_scalar:
        if isinstance(kw_sq, np.ndarray):
            ks = np.linalg.inv(kw_sq)
        else:
            ks = torch.linalg.inv(kw_sq).contiguous()
    elif kernel_type[-2:] == '_f' and not is_scalar:
        d = kw_sq.shape[-1]
        if isinstance(kw_sq, np.ndarray):
            cov_inv = np.linalg.inv(kw_sq)  # [... d d]
            w, v = np.linalg.eig(cov_inv)  # [... d], [... d d]
            w = w[..., None]**0.5 * np.eye(d).reshape(*([1]*(len(w.shape)-1)), d, d)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)
        else:
            cov_inv = torch.linalg.inv(kw_sq).contiguous()  # [... d d]
            w, v = torch.eig(cov_inv, eigenvectors=True)  # [... d], [... d d]
            w, v = w.real, v.real
            w = torch.diag_embed(w**0.5)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)  # [... d d]
    else:
        if kernel_type.startswith('nlin_'):
            ks = 1 / kw_sq**0.5
        else:
            ks = 1 / kw_sq
    return ks
//  [python]
def kw_sq_to_ks(kw_sq, kernel_type, is_scalar=False):
    if kernel_type[-2:] == '_a' and not is_scalar:
        if isinstance(kw_sq, np.ndarray):
            ks = np.linalg.inv(kw_sq)
        else:
            ks = torch.linalg.inv(kw_sq).contiguous()
    elif kernel_type[-2:] == '_f' and not is_scalar:
        d = kw_sq.shape[-1]
        if isinstance(kw_sq, np.ndarray):
            cov_inv = np.linalg.inv(kw_sq)  # [... d d]
            w, v = np.linalg.eig(cov_inv)  # [... d], [... d d]
            w = w[..., None]**0.5 * np.eye(d).reshape(*([1]*(len(w.shape)-1)), d, d)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)
        else:
            cov_inv = torch.linalg.inv(kw_sq).contiguous()  # [... d d]
            w, v = torch.linalg.eig(cov_inv)  # [... d], [... d d]
            w, v = w.real, v.real
            w = torch.diag_embed(w**0.5)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)  # [... d d]
    else:
        if kernel_type.startswith('nlin_'):
            ks = 1 / kw_sq**0.5
        else:
            ks = 1 / kw_sq
    return ks
//  [python]
def SPECTRE(U, temp_feats, n_dim, budget, oracle_clean_feats=None):

    projector = U[:, :n_dim].T # top left singular vectors
    temp_feats = torch.matmul(projector, temp_feats)

    if oracle_clean_feats is None:
        estimator = robust_estimation.BeingRobust(random_state=0, keep_filtered=True).fit((temp_feats.T).cpu().numpy())
        clean_mean = torch.FloatTensor(estimator.location_).cuda()
        filtered_feats = (torch.FloatTensor(estimator.filtered_).cuda() - clean_mean).T
        clean_covariance = torch.cov(filtered_feats)
    else:
        clean_feats = torch.matmul(projector, oracle_clean_feats)
        clean_covariance = torch.cov(clean_feats)
        clean_mean = clean_feats.mean(dim = 1)


    temp_feats = (temp_feats.T - clean_mean).T

    # whiten the data
    L, V = torch.linalg.eig(clean_covariance)
    L, V = L.real, V.real
    L = (torch.diag(L)**(1/2)+0.001).inverse()
    normalizer = torch.matmul(V, torch.matmul( L, V.T ) )
    temp_feats = torch.matmul(normalizer, temp_feats)

    # compute QUEscore
    taus = QUEscore(temp_feats, n_dim)

    sorted_indices = np.argsort(taus)
    n_samples = len(sorted_indices)

    budget = min(budget, n_samples//2) # default assumption : at least a half of samples in each class is clean

    suspicious = sorted_indices[-budget:]
    left = sorted_indices[:n_samples-budget]

    return suspicious, left
//  [python]
def SolveLayerEigensystem(omega, kx, ky, kp, ep2):
    """Solves the eigensystem.

    Args:
        omega (float): The angular frequency. Unused (from GRCWA).
        kx ():
        ky ():
        kp ():
        ep2 ():

    Returns:
        torch.tensor, torch.tensor: q, phi
    """
    nG = len(kx)

    k = torch.vstack((torch.diag(kx), torch.diag(ky)))
    kkT = torch_dot(k, torch_transpose(k))
    M = torch_dot(ep2, kp) - kkT

    q, phi = torch.eig(M)

    q = torch.sqrt(q)
    # branch cut choice
    q = torch.where(torch.imag(q) < 0.0, -q, q)
    return q, phi
//  [python]
def SolveLayerEigensystem(omega, kx, ky, kp, ep2):
    """Solves the eigensystem.

    Args:
        omega (float): The angular frequency. Unused (from GRCWA).
        kx ():
        ky ():
        kp ():
        ep2 ():

    Returns:
        torch.tensor, torch.tensor: q, phi
    """
    nG = len(kx)

    k = torch.vstack((torch.diag(kx), torch.diag(ky)))
    kkT = torch_dot(k, torch_transpose(k))
    M = torch_dot(ep2, kp) - kkT

    q, phi = torch.linalg.eig(M)

    q = torch.sqrt(q)
    # branch cut choice
    q = torch.where(torch.imag(q) < 0.0, -q, q)
    return q, phi
//  [python]
def eig(a: ArrayLike):
    a = _atleast_float_1(a)
    w, vt = torch.eig(a)

    if not a.is_complex() and w.is_complex() and (w.imag == 0).all():
        w = w.real
        vt = vt.real
    return w, vt
//  [python]
def eig(a: ArrayLike):
    a = _atleast_float_1(a)
    w, vt = torch.linalg.eig(a)

    if not a.is_complex() and w.is_complex() and (w.imag == 0).all():
        w = w.real
        vt = vt.real
    return w, vt
//  [python]
dist_mtx = syntheticCubeInfo["dist"]
LF_Norm = syntheticCubeInfo["LF"]
sourceFC = syntheticCubeInfo["Source FC"]
channelFC = syntheticCubeInfo["Channel FC"]

print(max(abs(torch.eig(Con_Mtx)[0].real)))
//  [python]
dist_mtx = syntheticCubeInfo["dist"]
LF_Norm = syntheticCubeInfo["LF"]
sourceFC = syntheticCubeInfo["Source FC"]
channelFC = syntheticCubeInfo["Channel FC"]

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

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

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

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

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

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

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

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

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

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

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

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

    return w, P, B, V
//  [python]
def eig(
    x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    result_tuple = NamedTuple(
        "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)]
    )
    eigenvalues, eigenvectors = torch.eig(x, out=out)
    return result_tuple(eigenvalues, eigenvectors)
//  [python]
def eig(
    x: torch.Tensor, /, *, out: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    result_tuple = NamedTuple(
        "eig", [("eigenvalues", torch.Tensor), ("eigenvectors", torch.Tensor)]
    )
    eigenvalues, eigenvectors = torch.linalg.eig(x, out=out)
    return result_tuple(eigenvalues, eigenvectors)
//  [python]
def whiten(self, x):
        cov = torch.matmul(x, x.T) / (x.shape[1] - 1)
        eig_vals, eig_vecs = torch.symeig(cov, eigenvectors=True)
        topk_indices = torch.topk(eig_vals.float().abs(), self.n_comp)[1]

        U = eig_vecs.float()
        lamb = eig_vals.float()[topk_indices].abs()
        lamb_inv_sqrt = torch.diag(1 / (torch.sqrt(lamb) + self.eps)).float()
        W = torch.matmul(lamb_inv_sqrt, U.T[topk_indices]).float()
        x_w = torch.matmul(W, x)
        return x_w, W
//  [python]
def whiten(self, x):
        cov = torch.matmul(x, x.T) / (x.shape[1] - 1)
        eig_vals, eig_vecs = torch.linalg.eig(cov)
        topk_indices = torch.topk(eig_vals.float().abs(), self.n_comp)[1]

        U = eig_vecs.float()
        lamb = eig_vals.float()[topk_indices].abs()
        lamb_inv_sqrt = torch.diag(1 / (torch.sqrt(lamb) + self.eps)).float()
        W = torch.matmul(lamb_inv_sqrt, U.T[topk_indices]).float()
        x_w = torch.matmul(W, x)
        return x_w, W
//  [python]
def get_min_max_real_eigenvalues(self) -> Tuple[np.float64, np.float64]:
        M = self.get_constraints()
        return (
            torch.min(torch.real(torch.eig(M)[0])).cpu().detach().numpy(),
            torch.max(torch.real(torch.eig(M)[0])).cpu().detach().numpy(),
        )
//  [python]
def get_min_max_real_eigenvalues(self) -> Tuple[np.float64, np.float64]:
        M = self.get_constraints()
        return (
            torch.min(torch.real(torch.linalg.eig(M)[0])).cpu().detach().numpy(),
            torch.max(torch.real(torch.linalg.eig(M)[0])).cpu().detach().numpy(),
        )
//  [python]
def _pca_torch(X, device, no_dims=50, verbose=False):
    if verbose:
        print("Preprocessing the data using PCA...")

    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    if not device == "cpu":
        with torch.cuda.device(device):
            (l, M) = torch.symeig(torch.mm(X.t(), X))
    else:
        (l, M) = torch.symeig(torch.mm(X.t(), X))
    # split M real
    for i in range(d):
        if l[i].imag != 0:
            M[:, i + 1] = M[:, i]
            i += 1

    M = M.real
    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def _pca_torch(X, device, no_dims=50, verbose=False):
    if verbose:
        print("Preprocessing the data using PCA...")

    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    if not device == "cpu":
        with torch.cuda.device(device):
            (l, M) = torch.linalg.eig(torch.mm(X.t(), X))
    else:
        (l, M) = torch.linalg.eig(torch.mm(X.t(), X))
    # split M real
    for i in range(d):
        if l[i].imag != 0:
            M[:, i + 1] = M[:, i]
            i += 1

    M = M.real
    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def fractional_matrix_power(matrix: Tensor, power: float, cache: Optional[Dict[str, Tensor]] = None):
    if cache is not None and "eigenvalues" in cache:
        complex_dtype = torch_complex_dtype_map[matrix.dtype]
        eigenvalues = cache["eigenvalues"].to(matrix.device, complex_dtype)
        eigenvectors = cache["eigenvectors"].to(matrix.device, complex_dtype)
        eigenvectors_inv = cache["eigenvectors_inv"].to(matrix.device, complex_dtype)
    else:
        eigenvalues, eigenvectors = torch.eig(matrix)
        eigenvectors_inv = torch.inverse(eigenvectors)
        if cache is not None:
            cache["eigenvalues"] = eigenvalues.to("cpu", torch.complex32)
            cache["eigenvectors"] = eigenvectors.to("cpu", torch.complex32)
            cache["eigenvectors_inv"] = eigenvectors_inv.to("cpu", torch.complex32)

    eigenvalues.pow_(power)
    result = eigenvectors @ torch.diag(eigenvalues) @ eigenvectors_inv
    return result.real.to(dtype=matrix.dtype)
//  [python]
def fractional_matrix_power(matrix: Tensor, power: float, cache: Optional[Dict[str, Tensor]] = None):
    if cache is not None and "eigenvalues" in cache:
        complex_dtype = torch_complex_dtype_map[matrix.dtype]
        eigenvalues = cache["eigenvalues"].to(matrix.device, complex_dtype)
        eigenvectors = cache["eigenvectors"].to(matrix.device, complex_dtype)
        eigenvectors_inv = cache["eigenvectors_inv"].to(matrix.device, complex_dtype)
    else:
        eigenvalues, eigenvectors = torch.linalg.eig(matrix)
        eigenvectors_inv = torch.linalg.inv(eigenvectors)
        if cache is not None:
            cache["eigenvalues"] = eigenvalues.to("cpu", torch.complex32)
            cache["eigenvectors"] = eigenvectors.to("cpu", torch.complex32)
            cache["eigenvectors_inv"] = eigenvectors_inv.to("cpu", torch.complex32)

    eigenvalues.pow_(power)
    result = eigenvectors @ torch.diag(eigenvalues) @ eigenvectors_inv
    return result.real.to(dtype=matrix.dtype)
//  [python]
def eig(a: Numeric, compute_eigvecs: bool = True):
    vals, vecs = torch.eig(a)
    return (vals, vecs) if compute_eigvecs else vals
//  [python]
def eig(a: Numeric, compute_eigvecs: bool = True):
    vals, vecs = torch.linalg.eig(a)
    return (vals, vecs) if compute_eigvecs else vals
//  [python]
def gradient_vis(gradients: List[torch.Tensor], **kwargs) -> Tuple[Figure, Axes]:
    default_figsize = (
        0.1,
        0.1,
    )
    figsize = kwargs.get("figsize", default_figsize)

    L, V = torch.eig(torch.stack(gradients)[:, 0, :])
    V_real = V.real

    fig, ax = plt.subplots(1, 1, figsize=figsize)
    scale = 2.0
    ax.scatter(V_real[..., 0, 1], V_real[..., 1, 1], c=range(V_real.shape[0], 0, -1), cmap="autumn", s=scale)
    ax.scatter(V_real[..., 0, 0], V_real[..., 1, 0], c=range(V_real.shape[0], 0, -1), cmap="winter", s=scale)

    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-1.1, 1.1)
    ax.axis("equal")

    return fig, ax
//  [python]
def gradient_vis(gradients: List[torch.Tensor], **kwargs) -> Tuple[Figure, Axes]:
    default_figsize = (
        0.1,
        0.1,
    )
    figsize = kwargs.get("figsize", default_figsize)

    L, V = torch.linalg.eig(torch.stack(gradients)[:, 0, :])
    V_real = V.real

    fig, ax = plt.subplots(1, 1, figsize=figsize)
    scale = 2.0
    ax.scatter(V_real[..., 0, 1], V_real[..., 1, 1], c=range(V_real.shape[0], 0, -1), cmap="autumn", s=scale)
    ax.scatter(V_real[..., 0, 0], V_real[..., 1, 0], c=range(V_real.shape[0], 0, -1), cmap="winter", s=scale)

    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-1.1, 1.1)
    ax.axis("equal")

    return fig, ax
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

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

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

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

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


    return w, p, p, B, V
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

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

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

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

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


    return w, p, p, B, V
//  [python]
def _get_steering_vector_evd(self, psd_s: torch.Tensor) -> torch.Tensor:
        r"""Estimate the steering vector by eigenvalue decomposition.

        Args:
            psd_s (torch.tensor): covariance matrix of speech
                Tensor of dimension `(..., freq, channel, channel)`

        Returns:
            torch.Tensor: the enhanced STFT
                Tensor of dimension `(..., freq, channel, 1)`
        """
        w, v = torch.linalg.eig(psd_s)  # (..., freq, channel, channel)
        _, indices = torch.max(w.abs(), dim=-1, keepdim=True)
        indices = indices.unsqueeze(-1)
        stv = v.gather(-1, indices.expand(psd_s.shape[:-1] + (1,)))  # (..., freq, channel, 1)
        return stv
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat

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

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

    # print("check", V @ torch.diag_embed(w) @ V.conj().transpose(-1, -2))

    # We require AP to be nearly skew-symmetric
    _A = AP + AP.transpose(-1, -2)
    if (
        err := torch.sum((_A - _A[0, 0] * torch.eye(N)) ** 2) / N
    ) > 1e-5:  # if not torch.allclose(_A - _A[0,0]*torch.eye(N), torch.zeros(N, N), atol=1e-5):
        print("WARNING: HiPPO matrix not skew symmetric", err)

    # Only keep half of each conjugate pair
    # w = w[..., 0::2].contiguous()
    # V = V[..., 0::2].contiguous()
    _, idx = torch.sort(w.imag)
    w_sorted = w[idx]
    V_sorted = V[:, idx]

    # There is an edge case when eigenvalues can be 0, which requires some machinery to handle
    # We use a huge hack here: Assume only one pair is 0, and that it is the first row/column of A (only happens in Fourier case)
    V = V_sorted[:, : N // 2]
    w = w_sorted[: N // 2]
    assert w[-2].abs() > 1e-4, "Only 1 zero eigenvalue allowed in diagonal part of A"
    if w[-1].abs() < 1e-4:
        V[:, -1] = 0.0
        V[0, -1] = 2**-0.5
        V[1, -1] = 2**-0.5 * 1j

    _AP = V @ torch.diag_embed(w) @ V.conj().transpose(-1, -2)
    # assert torch.allclose(2*_AP.real, AP, atol=1e-5)
    if (err := torch.sum((2 * _AP.real - AP) ** 2) / N) > 1e-5:
        print(
            "Warning: Diagonalization of A matrix not numerically precise - error", err
        )
    # print("check", V @ torch.diag_embed(w) @ V.conj().transpose(-1, -2))

    # # Override eigenvectors for 0 eigenvalues, to make them conjugate pairs
    # breakpoint()
    # rotate = torch.tensor([[1, 1], [1j, -1j]]) / 2**.5
    # # rotate = torch.tensor([[1, -1j], [1, 1j]]) / 2**.5
    # V_rot = (V.view(N, N//2, 2) @ rotate).view(N, N) # rotate every pair of eigenvectors
    # V = torch.where(w.repeat(N, 1) == 0, V_rot, V)

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

    C = initial_C(measure, N, dtype=dtype)
    B = contract("ij, j -> i", V_inv, B.to(V))  # V^* B
    C = contract("ij, j -> i", V_inv, C.to(V))  # V^* C
    P = contract("ij, ...j -> ...i", V_inv, P.to(V))  # V^* P

    return w, P, B, C, V
//  [python]
def _setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping."""
        self.dA, self.dB = self._setup_state()

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

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

        # Do special preprocessing for different step modes

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

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

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

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

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

        # Do special preprocessing for different step modes

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

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

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def compute_participation_ratio(x):
    """
    Calculates participation ratio of the correlation matrix of the input
    pr = (sum of eigenvalues)^2 / sum of (eigenvalues^2)
    :param x: (torch.tensor) input feature tensor (batch x features)
    :return: (torch.tensor)
             participation ratio (singleton tensor)
    """
    x_c = x - x.mean(dim=0)
    corr_matrix = torch.einsum('ij,kj->ik', x_c, x_c)
    eig_vals, _ = torch.linalg.eig(corr_matrix)
    eig_vals = eig_vals.real
    pr = eig_vals.sum() ** 2 / (eig_vals ** 2).sum()
    return pr
//  [python]
def participation_ratio(x):
    """
    Calculates participation ratio of the correlation matrix of the input
    pr = (sum of eigenvalues)^2 / sum of (eigenvalues^2)
    :param x: (torch.tensor) input feature tensor (batch x features)
    :return: (torch.tensor)
             participation ratio (singleton tensor)
    """
    x_c = x - x.mean(dim=0)

    corr_matrix = torch.einsum('ij,kj->ik', x_c, x_c)

    eig_vals, _ = torch.symeig(corr_matrix)

    pr = eig_vals.sum() ** 2 / (eig_vals ** 2).sum()

    return pr
//  [python]
def participation_ratio(x):
    """
    Calculates participation ratio of the correlation matrix of the input
    pr = (sum of eigenvalues)^2 / sum of (eigenvalues^2)
    :param x: (torch.tensor) input feature tensor (batch x features)
    :return: (torch.tensor)
             participation ratio (singleton tensor)
    """
    x_c = x - x.mean(dim=0)

    corr_matrix = torch.einsum('ij,kj->ik', x_c, x_c)

    eig_vals, _ = torch.linalg.eig(corr_matrix)
    eig_vals = eig_vals.real

    pr = eig_vals.sum() ** 2 / (eig_vals ** 2).sum()

    return pr
//  [python]
def eig(x):
    return torch.eig(x)
//  [python]
def eig(x):
    return torch.linalg.eig(x)
//  [python]
def _ncut_relabel(w, subgraph_nodes, thresh, num_cuts, original_labels, new_labels):
    d = torch.diag(w.sum(0))
    m = w.shape[0]
    if m > 2:
        d2 = copy.deepcopy(d)
        # Since d is diagonal, we can directly operate on its data
        # the inverse of the square root
        d2 = torch.diag(torch.reciprocal(torch.sqrt(torch.diag(d2))))
        # Refer Shi & Malik 2001, Equation 7, Page 891
        A = torch.matmul(torch.matmul(d2, (d - w)), d2)
        # v0 = torch.rand(A.shape[0])
        # vals, vectors = linalg.eigsh(A.cpu().numpy(), which='SM', v0=v0, k=min(100, m - 2))
        # vals, vectors = np.real(vals), np.real(vectors)
        # index2 = _ncut_cy.argmin2(vals)
        # ev = torch.from_numpy(vectors[:, index2]).to(w.device)

        k=min(100, m - 2)
        vals, vectors = torch.linalg.eig(A)
        vals, vectors = torch.real(vals), torch.real(vectors)
        vals, index = vals.sort()
        # vals = vals[:k]
        vectors = vectors[:, index[:k]]
        index2 = 1 if k >= 2 else 0
        ev = vectors[:, index2]
        
        cut_mask, mcut = get_min_ncut_torch(ev, d, w, num_cuts)
        if (mcut < thresh):
            sub1 = torch.where(cut_mask == True)[0]
            sub2 = torch.where(cut_mask != True)[0]
            _ncut_relabel(w[sub1][:, sub1], subgraph_nodes[sub1], thresh, num_cuts, original_labels, new_labels)
            _ncut_relabel(w[sub2][:, sub2], subgraph_nodes[sub2], thresh, num_cuts, original_labels, new_labels)
            return

    _label_all_torch(subgraph_nodes, original_labels, new_labels)
//  [python]
def get_log_like_gradient_and_hessian(self, X, eps=1e-2, compute_hessian=False):

        log_pi = self.forward(X)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.autograd.functional.hessian(self.forward, X)
            lamda, Q = torch.eig(H)
            lamda, Q = torch.real(lamda), torch.real(Q)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / lamda_prime
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def get_log_like_gradient_and_hessian(self, X, eps=1e-2, compute_hessian=False):

        log_pi = self.forward(X)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.autograd.functional.hessian(self.forward, X)
            lamda, Q = torch.linalg.eig(H)
            lamda, Q = torch.real(lamda), torch.real(Q)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / lamda_prime
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def rotation_from_matrix(R, eps=1e-4):
    w, W = torch.eig(R.T)
    i = torch.where(abs(torch.real(w) - 1.0) < eps)[0]
    if (len(i)==0):
        i = torch.tensor([0])
        print ('rotation_from_matrix w',torch.real(w))
        print ('rotation_from_matrix R.T',R.T)
    axis = torch.real(W[:, i[-1]]).squeeze()

    cosa = (torch.trace(R) - 1.0) / 2.0
    if abs(axis[2]) > eps:
        sina = (R[1, 0] + (cosa-1.0)*axis[0]*axis[1]) / axis[2]
    elif abs(axis[1]) > eps:
        sina = (R[0, 2] + (cosa-1.0)*axis[0]*axis[2]) / axis[1]
    else:
        sina = (R[2, 1] + (cosa-1.0)*axis[1]*axis[2]) / axis[0]
    angle = torch.atan2(sina, cosa)

    return angle, axis
//  [python]
def rotation_from_matrix(R, eps=1e-4):
    w, W = torch.linalg.eig(R.T)
    i = torch.where(abs(torch.real(w) - 1.0) < eps)[0]
    if (len(i)==0):
        i = torch.tensor([0])
        print ('rotation_from_matrix w',torch.real(w))
        print ('rotation_from_matrix R.T',R.T)
    axis = torch.real(W[:, i[-1]]).squeeze()

    cosa = (torch.trace(R) - 1.0) / 2.0
    if abs(axis[2]) > eps:
        sina = (R[1, 0] + (cosa-1.0)*axis[0]*axis[1]) / axis[2]
    elif abs(axis[1]) > eps:
        sina = (R[0, 2] + (cosa-1.0)*axis[0]*axis[2]) / axis[1]
    else:
        sina = (R[2, 1] + (cosa-1.0)*axis[1]*axis[2]) / axis[0]
    angle = torch.atan2(sina, cosa)

    return angle, axis
//  [python]
def _setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping."""
        self.dA, self.dB = self._setup_state()

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

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

        # Do special preprocessing for different step modes

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

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

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

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

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

        # Do special preprocessing for different step modes

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

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

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "DPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def forward(self, H1, H2):
        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()

        o1 = H1.size(0)
        o2 = H2.size(0)

        m = H1.size(1)

        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=H1.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

        # Workaround !!! USE BATCH > 16
        # [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
        # [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

        # Newest way but not debugged
        # [D1, V1] = torch.linalg.eig(SigmaHat11)
        # [D2, V2] = torch.linalg.eig(SigmaHat22)

        # Original
        [D1, V1] = torch.eig(SigmaHat11, eigenvectors=True)
        [D2, V2] = torch.eig(SigmaHat22, eigenvectors=True)

        D1 = D1.unsqueeze(1)
        posInd1 = torch.gt(abs(D1[:, 0]), eps).nonzero()[:, 0]
        D1 = D1[posInd1, 0]
        V1 = V1[:, posInd1]
        D1 = torch.squeeze(D1)  # Remove extra dimensions from D1
        # Reshape D1 to 1-dimensional tensor
        D2 = D2.unsqueeze(1)
        # posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        posInd2 = torch.gt(abs(D2[:, 0]), eps).nonzero()[:, 0]
        D2 = D2[posInd2, 0]
        V2 = V2[:, posInd2]
        D2 = torch.squeeze(D2)

        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12.to(torch.complex64)), SigmaHat22RootInv)

        if self.use_all_singular_values:

            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))

        else:

            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0]) * r1).to(H1.device))
            # U, V = torch.symeig(trace_TT, eigenvectors=True)
            U, V = torch.linalg.eigh(trace_TT)
            U = torch.where(U > eps, U, (torch.ones(U.shape).float() * eps).to(H1.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def forward(self, H1, H2):
        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()

        o1 = H1.size(0)
        o2 = H2.size(0)

        m = H1.size(1)

        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=H1.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

        # Workaround !!! USE BATCH > 16
        # [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
        # [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

        # Newest way but not debugged
        [D1, V1] = torch.linalg.eig(SigmaHat11)
        [D2, V2] = torch.linalg.eig(SigmaHat22)

        # Original
        # [D1, V1] = torch.eig(SigmaHat11, eigenvectors=True)
        # [D2, V2] = torch.eig(SigmaHat22, eigenvectors=True)

        D1 = D1.unsqueeze(1)
        posInd1 = torch.gt(abs(D1[:, 0]), eps).nonzero()[:, 0]
        D1 = D1[posInd1, 0]
        V1 = V1[:, posInd1]
        D1 = torch.squeeze(D1)  # Remove extra dimensions from D1
        # Reshape D1 to 1-dimensional tensor
        D2 = D2.unsqueeze(1)
        # posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        posInd2 = torch.gt(abs(D2[:, 0]), eps).nonzero()[:, 0]
        D2 = D2[posInd2, 0]
        V2 = V2[:, posInd2]
        D2 = torch.squeeze(D2)

        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12.to(torch.complex64)), SigmaHat22RootInv)

        if self.use_all_singular_values:

            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))

        else:

            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0]) * r1).to(H1.device))
            # U, V = torch.symeig(trace_TT, eigenvectors=True)
            U, V = torch.linalg.eigh(trace_TT)
            U = torch.where(U > eps, U, (torch.ones(U.shape).float() * eps).to(H1.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def forward(ctx, A):
        evalue, evec = torch.eig(A, eigenvectors=True)

        evecs = torch.view_as_real(evec)[..., 0]
        evals = torch.view_as_real(evalue)[..., 0]
        evalsnew = evals.clone()
        evecsnew = evecs.clone()

        if torch.isnan(evals).any():
            print(evals[evals != evals])

        xx = torch.argmax(torch.abs(evecs[:, 0, :]), dim=1, keepdim=True)
        xxx = xx.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)
        yy = torch.argmax(torch.abs(evecs[:, 1, :]), dim=1, keepdim=True)
        yyy = yy.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)
        zz = torch.argmax(torch.abs(evecs[:, 2, :]), dim=1, keepdim=True)
        zzz = zz.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)

        evalsnew[:, 0:1] = torch.gather(evals, 1, xx)
        evalsnew[:, 1:2] = torch.gather(evals, 1, yy)
        evalsnew[:, 2:3] = torch.gather(evals, 1, zz)
        evecsnew[:, :, 0:1] = torch.gather(evecs, 2, xxx)
        evecsnew[:, :, 1:2] = torch.gather(evecs, 2, yyy)
        evecsnew[:, :, 2:3] = torch.gather(evecs, 2, zzz)
        ctx.save_for_backward(evecsnew)

        return evalsnew
//  [python]
def forward(ctx, A):
        '''
        A is the hessian matrix but for all points of all patches, so shape is [batch_dim*x_dim*y_dim*z_dim,3,3]
        so I do evecs[:,0,:] to select for all the voxels, the line that correspond to the x direction (1 for y and 2 for z),
        then I extract the position where this direction is the biggest: xx = torch.argmax(torch.abs(evecs[:,0,:]),dim=1,keepdim=True)

        With torch.linalg.eig the normalized eigenvectors are given per column with respect to the eigenvalue and per row
        with respect to the directions, e.g. for a voxel:
        eval = [-256 -41 0]
        evec = [0.9 0.3 0; 0 0 1; -0.3 0.9 0]

        so eval[0]=-256 and its evec is [0.9; 0; -0.3] = [x; y; z]

        The xx and xxx step is done because the use of permutation matrix (as described in the artcile/thesis) was slower
        and so I just rearranged the eigenvalues and eigenvectors to be [most x-oriented, most y-oriented, most z-oriented].

        So for the example, we will have xx=0, yy=2, zz = 1 and eval will be rearranged as eval = [-256 0 -41]
        and xxx, yyy and zzz simply consist in repeating the values and reorganizing the eigenvector matrix which is used
        for the backpropagation, which will therefore be

        evec = [0.9 0 0.3; 0 1 0; -0.3 0 0.9]
        '''

        evalue, evec = torch.linalg.eig(A)

        evecs = torch.view_as_real(evec)[..., 0]
        evals = torch.view_as_real(evalue)[..., 0]
        evalsnew = evals.clone()
        evecsnew = evecs.clone()

        if torch.isnan(evals).any():
            print(evals[evals != evals])

        xx = torch.argmax(torch.abs(evecs[:, 0, :]), dim=1, keepdim=True)
        xxx = xx.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)
        yy = torch.argmax(torch.abs(evecs[:, 1, :]), dim=1, keepdim=True)
        yyy = yy.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)
        zz = torch.argmax(torch.abs(evecs[:, 2, :]), dim=1, keepdim=True)
        zzz = zz.repeat(1, evecs.size()[-1]).view(evecs.size()[0], evecs.size()[-1], 1)

        evalsnew[:, 0:1] = torch.gather(evals, 1, xx)
        evalsnew[:, 1:2] = torch.gather(evals, 1, yy)
        evalsnew[:, 2:3] = torch.gather(evals, 1, zz)
        evecsnew[:, :, 0:1] = torch.gather(evecs, 2, xxx)
        evecsnew[:, :, 1:2] = torch.gather(evecs, 2, yyy)
        evecsnew[:, :, 2:3] = torch.gather(evecs, 2, zzz)
        ctx.save_for_backward(evecsnew)

        return evalsnew
//  [python]
def eigenvalues(self) -> TT[TH.leading_dims : ..., TH.mdim]:
        """Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k."""
        return torch.eig(self.BA)[0]
//  [python]
def eigenvalues(self) -> TT[TH.leading_dims : ..., TH.mdim]:
        """Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k."""
        return torch.linalg.eig(self.BA).eigenvalues
//  [python]
def pca_denoised(im,denoisers, sigmas, B, N):

    Ls={}
    Qs={}
    denoised_ims_centered = {}
    clean_ims = torch.tile(im , dims = (B,1,1,1)).to(device)
    noises = torch.randn_like(clean_ims, device=device)

    for sigma in sigmas:
        noisy_ims = clean_ims +  noises* sigma/255

        with torch.no_grad():
            residuals = denoisers[N](noisy_ims )

        denoised_ims = noisy_ims - residuals
        denoised_ims_centered[sigma] = denoised_ims- denoised_ims.mean(dim= 0, keepdims=True)
        cov = torch.cov(torch.flatten(denoised_ims_centered[sigma].squeeze(), start_dim= 1, end_dim=2).T)
        L, Q = torch.linalg.eig(cov)
        Ls[sigma] = torch.real(L)
        Qs[sigma] = torch.real(Q)

    return Ls, Qs, denoised_ims