jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def proj_spectrahedron_torch(rho, device, map_method, P_proj=2, trace_flag=1):
    """
    Select the state-mapping method according to the given parameters ``map_method`` and ``P_proj``.

    Args:
        rho (tensor): Matrix that does not satisfy the density matrix property.
        device (torch.device): GPU or CPU.
        map_method (str): State-mapping method, include ['proj_F', 'proj_S', 'proj_A'].
        P_proj (float): P order.
        trace_flag (1, not 1): 1 is required for the density matrix to satisfy a diagonal 
            sum of one, otherwise it is not required.

    Returns:
        The real density matrix.
    """
    eigenvalues, eigenvecs = torch.symeig(rho, eigenvectors=True)  # eigenvalues[i], eigenvecs[:, i]

    #eigenvalues = softmax(eigenvalues, 0)
    if map_method == 'proj_F':
        eigenvalues = eigenvalues_trans_F(eigenvalues, device)
    elif map_method == 'proj_S':
        eigenvalues = eigenvalues_trans_S(eigenvalues, device)
    elif map_method == 'proj_A':
        eigenvalues = eigenvalues_trans_abs(eigenvalues, P_proj)
    else:
        print('we have not this map method! please check setting!!!')

    A = eigenvecs * eigenvalues
    rho = torch.matmul(A, eigenvecs.T.conj())
    rho = 0.5 * (rho + rho.T.conj())
    if trace_flag == 1:
        rho /= torch.trace(rho)  # prevent errors caused by computing accuracy
    return rho
//  [python]
def proj_spectrahedron_torch(rho, device, map_method, P_proj=2, trace_flag=1):
    """
    Select the state-mapping method according to the given parameters ``map_method`` and ``P_proj``.

    Args:
        rho (tensor): Matrix that does not satisfy the density matrix property.
        device (torch.device): GPU or CPU.
        map_method (str): State-mapping method, include ['proj_F', 'proj_S', 'proj_A'].
        P_proj (float): P order.
        trace_flag (1, not 1): 1 is required for the density matrix to satisfy a diagonal 
            sum of one, otherwise it is not required.

    Returns:
        The real density matrix.
    """
    eigenvalues, eigenvecs = torch.linalg.eigh(rho)  # eigenvalues[i], eigenvecs[:, i]

    #eigenvalues = softmax(eigenvalues, 0)
    if map_method == 'proj_F':
        eigenvalues = eigenvalues_trans_F(eigenvalues, device)
    elif map_method == 'proj_S':
        eigenvalues = eigenvalues_trans_S(eigenvalues, device)
    elif map_method == 'proj_A':
        eigenvalues = eigenvalues_trans_abs(eigenvalues, P_proj)
    else:
        print('we have not this map method! please check setting!!!')

    A = eigenvecs * eigenvalues
    rho = torch.matmul(A, eigenvecs.T.conj())
    rho = 0.5 * (rho + rho.T.conj())
    if trace_flag == 1:
        rho /= torch.trace(rho)  # prevent errors caused by computing accuracy
    return rho
//  [python]
def collate_pad(batch):
    E = []
    U = []
    X = []
    F = []
    Y = []
    
    max_nodes = min(max([data.num_nodes for data in batch]), 128)

    for data in batch:
        length = data.num_nodes
        e = data.e
        u = data.u.view(length, length)
        x = data.x
        f = data.edge_attr

        if length > max_nodes:
            src, dst = data.edge_index
            A = torch.zeros([length, length], dtype=torch.float)
            A[src, dst] = 1.0
            A = A[:max_nodes, :max_nodes]
            deg = torch.sum(A, axis=0).squeeze()
            deg = torch.clamp(deg, min=1.0) ** -0.5
            D = torch.diag(deg)
            A_ = D @ A @ D

            pad_e, pad_u = torch.eig(A_, eigenvectors=True)
            pad_x = x[:max_nodes, :] + 1

            fdim = f.size(-1)
            pad_f = torch.zeros([length, length, fdim], dtype=torch.long)
            pad_f[src, dst] = f + 1
            pad_f = pad_f[:max_nodes, :max_nodes] 
        else:
            pad_e = e.new_zeros([max_nodes])
            pad_e[:length] = e

            pad_u = u.new_zeros([max_nodes, max_nodes])
            pad_u[:length, :length] = u

            xdim = x.size(-1)
            pad_x = x.new_zeros([max_nodes, xdim])
            pad_x[:length, :] = x + 1

            fdim = f.size(-1)
            src, dst = data.edge_index
            pad_f = f.new_zeros([max_nodes, max_nodes, fdim])
            pad_f[src, dst, :] = f + 1 

        E.append(pad_e)
        U.append(pad_u)
        X.append(pad_x)
        F.append(pad_f)
        Y.append(data.y.squeeze())

    return torch.stack(E, 0), torch.stack(U, 0), torch.stack(X, 0), torch.stack(F, 0), torch.stack(Y, 0)
//  [python]
def collate_pad(batch):
    E = []
    U = []
    X = []
    F = []
    Y = []
    
    max_nodes = min(max([data.num_nodes for data in batch]), 128)

    for data in batch:
        length = data.num_nodes
        e = data.e
        u = data.u.view(length, length)
        x = data.x
        f = data.edge_attr

        if length > max_nodes:
            src, dst = data.edge_index
            A = torch.zeros([length, length], dtype=torch.float)
            A[src, dst] = 1.0
            A = A[:max_nodes, :max_nodes]
            deg = torch.sum(A, axis=0).squeeze()
            deg = torch.clamp(deg, min=1.0) ** -0.5
            D = torch.diag(deg)
            A_ = D @ A @ D

            pad_e, pad_u = torch.linalg.eigh(A_)
            pad_x = x[:max_nodes, :] + 1

            fdim = f.size(-1)
            pad_f = torch.zeros([length, length, fdim], dtype=torch.long)
            pad_f[src, dst] = f + 1
            pad_f = pad_f[:max_nodes, :max_nodes] 
        else:
            pad_e = e.new_zeros([max_nodes])
            pad_e[:length] = e

            pad_u = u.new_zeros([max_nodes, max_nodes])
            pad_u[:length, :length] = u

            xdim = x.size(-1)
            pad_x = x.new_zeros([max_nodes, xdim])
            pad_x[:length, :] = x + 1

            fdim = f.size(-1)
            src, dst = data.edge_index
            pad_f = f.new_zeros([max_nodes, max_nodes, fdim])
            pad_f[src, dst, :] = f + 1 

        E.append(pad_e)
        U.append(pad_u)
        X.append(pad_x)
        F.append(pad_f)
        Y.append(data.y.squeeze())

    return torch.stack(E, 0), torch.stack(U, 0), torch.stack(X, 0), torch.stack(F, 0), torch.stack(Y, 0)
//  [python]
def _symapply(x, f, *, wmin=None, wmax=None):
    r"""Template function acting on stacked symmetric matrices that applies a
    given analytic function on them via eigenvalue decomposition.
    """
    # try:
    w, v = torch.symeig(x)
    # except RuntimeError as e:
    #    print(e)
    #     torch.save(x, "x.pt")
    #   sys.exit(0)

    # print("W", w[0])
    # print("V", v[0])
    if wmin is not None or wmax is not None:
        w.data.clamp_(min=wmin, max=wmax)

    return _mvmt(v, f(w), v)
//  [python]
def _symapply(x, f, *, wmin=None, wmax=None):
    r"""Template function acting on stacked symmetric matrices that applies a
    given analytic function on them via eigenvalue decomposition.
    """
    # try:
    w, v = torch.linalg.eigh(x)
    # except RuntimeError as e:
    #    print(e)
    #     torch.save(x, "x.pt")
    #   sys.exit(0)

    # print("W", w[0])
    # print("V", v[0])
    if wmin is not None or wmax is not None:
        w.data.clamp_(min=wmin, max=wmax)

    return _mvmt(v, f(w), v)
//  [python]
def _compute_eigenspace(self):
        self.eigenvalues.data, self.eigenvectors.data = torch.symeig(
            self._compute_autorcorrelation(), eigenvectors=True, upper=True
        )  #.type(self.data_dtype)
        self.eigenvalues.data, idx = self.eigenvalues.sort(descending=True)
        # correct numerical error, matrix must be positivly semi-definitie
        self.eigenvalues[self.eigenvalues < 0] = 0
        self.eigenvectors.data = self.eigenvectors[:, idx]
//  [python]
def _compute_eigenspace(self):
        self.eigenvalues.data, self.eigenvectors.data = torch.linalg.eigh(
            self._compute_autorcorrelation(),  UPLO='U'
        )  #.type(self.data_dtype)
        self.eigenvalues.data, idx = self.eigenvalues.sort(descending=True)
        # correct numerical error, matrix must be positivly semi-definitie
        self.eigenvalues[self.eigenvalues < 0] = 0
        self.eigenvectors.data = self.eigenvectors[:, idx]
//  [python]
def from_distribution(self, dist="gaussian", is_symmetric=True):
        """

        Args:
            dist (str, optional): Distribution of eigenvalues.  Defaults to "gaussian".
            is_symmetric (bool, optional): If true, sampled matrices will be symmetric. Defaults to True.
        """
        # Sample new eigenvalues
        if self.dist == "gaussian":
            x = self.sigma * torch.randn(self.N, self.d, 1)
        elif self.dist == "uniform":
            x = (
                -2 * math.sqrt(self.d) * 10 * torch.rand(self.N, self.d, 1)
                + math.sqrt(self.d) * 10
            )
        elif self.dist == "laplace":
            m = torch.distributions.Laplace(
                torch.tensor([0.0]), torch.tensor([self.sigma / math.sqrt(2)])
            )
            x = m.rsample(sample_shape=torch.Size([self.N, self.d]))
        diag = torch.eye(x.shape[1]) * x[:, None]

        M = torch.randn(self.N, 1, self.d, self.d)
        if is_symmetric:
            # Create symmetric matrices
            M = torch.triu(M, 0) + torch.transpose(torch.triu(M, 1), 2, 3)
            # Compute eigenvectors of hermitian matrix
            P = torch.symeig(M,eigenvectors=True)[1]
            # Finally, construct the resulting matrix batch with specified eigenvalues
            self.X = torch.matmul(torch.matmul(P, diag), torch.transpose(P, 2, 3))
        else:
            self.X = torch.matmul(torch.matmul(M, diag), torch.linalg.inv(M))
//  [python]
def from_distribution(self, dist="gaussian", is_symmetric=True):
        """

        Args:
            dist (str, optional): Distribution of eigenvalues.  Defaults to "gaussian".
            is_symmetric (bool, optional): If true, sampled matrices will be symmetric. Defaults to True.
        """
        # Sample new eigenvalues
        if self.dist == "gaussian":
            x = self.sigma * torch.randn(self.N, self.d, 1)
        elif self.dist == "uniform":
            x = (
                -2 * math.sqrt(self.d) * 10 * torch.rand(self.N, self.d, 1)
                + math.sqrt(self.d) * 10
            )
        elif self.dist == "laplace":
            m = torch.distributions.Laplace(
                torch.tensor([0.0]), torch.tensor([self.sigma / math.sqrt(2)])
            )
            x = m.rsample(sample_shape=torch.Size([self.N, self.d]))
        diag = torch.eye(x.shape[1]) * x[:, None]

        M = torch.randn(self.N, 1, self.d, self.d)
        if is_symmetric:
            # Create symmetric matrices
            M = torch.triu(M, 0) + torch.transpose(torch.triu(M, 1), 2, 3)
            # Compute eigenvectors of hermitian matrix
            P = torch.linalg.eigh(M)[1]
            # Finally, construct the resulting matrix batch with specified eigenvalues
            self.X = torch.matmul(torch.matmul(P, diag), torch.transpose(P, 2, 3))
        else:
            self.X = torch.matmul(torch.matmul(M, diag), torch.linalg.inv(M))
//  [python]
def _fit_plane(x):
    x = x - x.mean(dim=0)[None]
    cov = (x.t() @ x) / x.shape[0]
    _, e_vec = torch.symeig(cov)
    return e_vec
//  [python]
def _fit_plane(x):
    x = x - x.mean(dim=0)[None]
    cov = (x.t() @ x) / x.shape[0]
    _, e_vec = torch.linalg.eigh(cov)
    return e_vec
//  [python]
def RtoP(R):
    # return P = R^(-0.25)
    L, U = torch.symeig(R)
    P = U @ torch.diag(L ** (-0.25)) @ U.t()
    return P
//  [python]
def RtoP(R):
    # return P = R^(-0.25)
    L, U = torch.linalg.eigh(R)
    P = U @ torch.diag(L ** (-0.25)) @ U.t()
    return P
//  [python]
def compute_eigendecomposition(A, k=None, eps=1e-8):
    """Eigendecomposition of a square matrix A.

    Args:
        A: square matrix A
        k: number of eigenvectors
        eps: small error term

    Returns:
        evals (k): eigenvalues of the Laplacian
        evecs (V,k): matrix of eigenvectors of the Laplacian
    """
    if A is None:
        return None

    if k is None:
        A = A.to_dense()
    else:
        indices, values, size = A.indices(), A.values(), A.size()
        A = sp.coo_array((values, (indices[0], indices[1])), shape=size)

    failcount = 0
    while True:
        try:
            if k is None:
                evals, evecs = torch.symeig(A, eigenvectors=True)
            else:
                evals, evecs = sp.linalg.eigsh(A, k=k, which="SM")
                evals, evecs = torch.tensor(evals), torch.tensor(evecs)

            evals = torch.clamp(evals, min=0.0)
            evecs *= np.sqrt(len(evecs))

            break
        except Exception as e:  # pylint: disable=broad-exception-caught
            print(e)
            if failcount > 3:
                raise ValueError("failed to compute eigendecomp") from e
            failcount += 1
            print("--- decomp failed; adding eps ===> count: " + str(failcount))
            A += sp.eye(A.shape[0]) * (eps * 10 ** (failcount - 1))

    return evals, evecs
//  [python]
def compute_eigendecomposition(A, k=None, eps=1e-8):
    """Eigendecomposition of a square matrix A.

    Args:
        A: square matrix A
        k: number of eigenvectors
        eps: small error term

    Returns:
        evals (k): eigenvalues of the Laplacian
        evecs (V,k): matrix of eigenvectors of the Laplacian
    """
    if A is None:
        return None

    if k is None:
        A = A.to_dense()
    else:
        indices, values, size = A.indices(), A.values(), A.size()
        A = sp.coo_array((values, (indices[0], indices[1])), shape=size)

    failcount = 0
    while True:
        try:
            if k is None:
                evals, evecs = torch.linalg.eigh(A)  # pylint: disable=not-callable
            else:
                evals, evecs = sp.linalg.eigsh(A, k=k, which="SM")
                evals, evecs = torch.tensor(evals), torch.tensor(evecs)

            evals = torch.clamp(evals, min=0.0)
            evecs *= np.sqrt(len(evecs))

            break
        except Exception as e:  # pylint: disable=broad-exception-caught
            print(e)
            if failcount > 3:
                raise ValueError("failed to compute eigendecomp") from e
            failcount += 1
            print("--- decomp failed; adding eps ===> count: " + str(failcount))
            A += sp.eye(A.shape[0]) * (eps * 10 ** (failcount - 1))

    return evals, evecs
//  [python]
def _sqrt_Nystrom(K:Tensor, mask:Tensor) -> Tensor:
    """
    Compute the Nystrom approx square root `Ny(K)` of a kernel matrix.
    A smooth function is used for near-zero eigenvalues.
    """
    D, V = torch.symeig(K[mask])
    D_sqrt = torch.sqrt(torch.clamp(D, 0))
    D_sqrt_inv = D_sqrt / torch.clamp(D, D[-1] * 1e-4)
    Q = torch.zeros_like(K)
    Q[mask] = V * D_sqrt.view((1, -1))
    Q[~mask] = K[~mask] @ V * D_sqrt_inv.view((1, -1))
    return Q
//  [python]
def _sqrt_Nystrom(K:Tensor, mask:Tensor) -> Tensor:
    """
    Compute the Nystrom approx square root `Ny(K)` of a kernel matrix.
    A smooth function is used for near-zero eigenvalues.
    """
    D, V = torch.linalg.eigh(K[mask])
    D_sqrt = torch.sqrt(torch.clamp(D, 0))
    D_sqrt_inv = D_sqrt / torch.clamp(D, D[-1] * 1e-4)
    Q = torch.zeros_like(K)
    Q[mask] = V * D_sqrt.view((1, -1))
    Q[~mask] = K[~mask] @ V * D_sqrt_inv.view((1, -1))
    return Q
//  [python]
def eigh(x):
    return t.linalg.eigh(x)
//  [python]
def eigh(x):
    return torch.linalg.eigh(x)
//  [python]
def demo_degenerate_eigh_numerical_instable():
    np_rng = np.random.default_rng(233)
    np1 = np_rng.normal(size=(4,4))
    EVC_ = np.kron(np.eye(2),np.linalg.svd(np_rng.normal(size=(2,2)))[0])
    eps_list = np.concatenate([np.zeros(1), 10**np.linspace(-18, -13, 10)]) #about 1e-16
    z0 = []
    for eps in eps_list:
        tmp0 = np.array([1, 2, 2+eps, 3])
        np0 = (EVC_ * tmp0) @ EVC_.T
        torch0 = torch.tensor(np0, dtype=torch.float64, requires_grad=True)
        torch1 = torch.tensor(np1, dtype=torch.float64)
        EVL,EVC = torch.symeig((torch0+torch0.T)/2)
        loss = torch.trace(((EVC*torch.log(EVL)) @ EVC.T) @ torch1)
        loss.backward()
        z0.append(torch0.grad.detach().numpy().copy())
        if np.any(np.isnan(z0[-1])):
            print(eps)
    z0 = np.stack(z0)
//  [python]
def demo_degenerate_eigh_numerical_instable():
    np_rng = np.random.default_rng(233)
    np1 = np_rng.normal(size=(4,4))
    EVC_ = np.kron(np.eye(2),np.linalg.svd(np_rng.normal(size=(2,2)))[0])
    eps_list = np.concatenate([np.zeros(1), 10**np.linspace(-18, -13, 10)]) #about 1e-16
    z0 = []
    for eps in eps_list:
        tmp0 = np.array([1, 2, 2+eps, 3])
        np0 = (EVC_ * tmp0) @ EVC_.T
        torch0 = torch.tensor(np0, dtype=torch.float64, requires_grad=True)
        torch1 = torch.tensor(np1, dtype=torch.float64)
        EVL,EVC = torch.linalg.eigh((torch0+torch0.T)/2)
        loss = torch.trace(((EVC*torch.log(EVL)) @ EVC.T) @ torch1)
        loss.backward()
        z0.append(torch0.grad.detach().numpy().copy())
        if np.any(np.isnan(z0[-1])):
            print(eps)
    z0 = np.stack(z0)
//  [python]
def fit(self, n_components=None, whiten=False, method='e-decomposition'):
        if n_components is None:
            n_components = self.cov_mat.shape[0]
        self.whiten = whiten
        if method == 'e-decomposition':
            evals, Q = torch.symeig(self.cov_mat)
            idx = evals.argsort().flip(0)
            self._evals = evals[idx]
            evals = evals[idx][:n_components]
            self.Q = Q[:,idx]
            self.Q = self.Q[:, :n_components]
            self._Q = self.Q.clone()
            self.Lambda = torch.diag(evals)
            if self.whiten: self.Q = self.Q@torch.diag(evals**-0.5)
        elif method == 'SVD':
            u, s, vh = torch.linalg.svd(self.dstar,full_matrices=False)
            evals = (1/(self.dstar.shape[-1] - 1))*(s**2)
            self._evals = evals
            evals = evals[:n_components]
            self.Q = u[:, :n_components]
            self._Q = self.Q.clone()
            self.Lambda = torch.diag(evals)
            if self.whiten: self.Q = self.Q@torch.diag(evals**-0.5)
//  [python]
def fit(self, n_components=None, whiten=False, method='e-decomposition'):
        if n_components is None:
            n_components = self.cov_mat.shape[0]
        self.whiten = whiten
        if method == 'e-decomposition':
            evals, Q = torch.linalg.eigh(self.cov_mat)
            idx = evals.argsort().flip(0)
            self._evals = evals[idx]
            evals = evals[idx][:n_components]
            self.Q = Q[:,idx]
            self.Q = self.Q[:, :n_components]
            self._Q = self.Q.clone()
            self.Lambda = torch.diag(evals)
            if self.whiten: self.Q = self.Q@torch.diag(evals**-0.5)
        elif method == 'SVD':
            u, s, vh = torch.linalg.svd(self.dstar,full_matrices=False)
            evals = (1/(self.dstar.shape[-1] - 1))*(s**2)
            self._evals = evals
            evals = evals[:n_components]
            self.Q = u[:, :n_components]
            self._Q = self.Q.clone()
            self.Lambda = torch.diag(evals)
            if self.whiten: self.Q = self.Q@torch.diag(evals**-0.5)
//  [python]
def from_coords(
        cls,
        coords_x: torch.Tensor,
        bond_idxs: torch.Tensor,
        constr: ConstraintDict | None = None,
    ) -> "DLC":
        """Projects a set of cartesian coordinates onto a reduced set of delocalized
        internal coordinates.

        Args:
            coords_x: The coordinates with shape=(n_atoms, 3).
            bond_idxs: The atoms involved in each bond with shape=(n_bonds, 2).
            constr: A dictionary of constraints on the internal coordinates. This
                dictionary should be of the form ``constr[ic_type] = (idxs, values)``
                where ``idxs`` is a tensor of the indices of the atoms involved in the
                internal coordinate and ``values`` is a tensor of the values to
                constrain the internal coordinate at.

        Returns:
            The internal coordinate representation.
        """
        coords_x = coords_x.double()

        ric = RIC.from_coords(coords_x, bond_idxs)

        b_matrix = ric.compute_b(coords_x)
        g_matrix = b_matrix @ b_matrix.T

        eig_vals, eig_vecs = torch.symeig(g_matrix)
        eig_mask = torch.abs(eig_vals) > 1e-6

        v = eig_vecs[:, eig_mask]

        if constr is not None and len(constr) == 0:
            constr = None
        if constr is not None:
            v = cls._project_constr(v, constr, ric)

        return cls(ric.idxs, v, constr)
//  [python]
def from_coords(
        cls,
        coords_x: torch.Tensor,
        bond_idxs: torch.Tensor,
        constr: ConstraintDict | None = None,
    ) -> "DLC":
        """Projects a set of cartesian coordinates onto a reduced set of delocalized
        internal coordinates.

        Args:
            coords_x: The coordinates with shape=(n_atoms, 3).
            bond_idxs: The atoms involved in each bond with shape=(n_bonds, 2).
            constr: A dictionary of constraints on the internal coordinates. This
                dictionary should be of the form ``constr[ic_type] = (idxs, values)``
                where ``idxs`` is a tensor of the indices of the atoms involved in the
                internal coordinate and ``values`` is a tensor of the values to
                constrain the internal coordinate at.

        Returns:
            The internal coordinate representation.
        """
        coords_x = coords_x.double()

        ric = RIC.from_coords(coords_x, bond_idxs)

        b_matrix = ric.compute_b(coords_x)
        g_matrix = b_matrix @ b_matrix.T

        eig_vals, eig_vecs = torch.linalg.eigh(g_matrix)
        eig_mask = torch.abs(eig_vals) > 1e-6

        v = eig_vecs[:, eig_mask]

        if constr is not None and len(constr) == 0:
            constr = None
        if constr is not None:
            v = cls._project_constr(v, constr, ric)

        return cls(ric.idxs, v, constr)
//  [python]
def compute_features(self, noisy_data):
        E_t = noisy_data["E_t"]
        mask = noisy_data["node_mask"]
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = self.compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        # debug for protein dataset
        if L.isnan().any():
            import pdb; pdb.set_trace()
        eigvals, eigvectors = torch.symeig(L, eigenvectors=True)
        eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)
        eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
        # Retrieve eigenvalues features
        n_connected_comp, batch_eigenvalues = self.eigenvalues_features(
            eigenvalues=eigvals, num_eigenvalues=self.num_eigenvalues
        )
        # Retrieve eigenvectors features
        evector_feat = self.eigenvector_features(
            vectors=eigvectors,
            node_mask=noisy_data["node_mask"],
            n_connected=n_connected_comp,
            num_eigenvectors=self.num_eigenvectors,
        )

        evalue_feat = torch.hstack((n_connected_comp, batch_eigenvalues))
        return evalue_feat, evector_feat
//  [python]
def compute_features(self, noisy_data):
        E_t = noisy_data["E_t"]
        mask = noisy_data["node_mask"]
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = self.compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        # debug for protein dataset
        if L.isnan().any():
            import pdb; pdb.set_trace()
        eigvals, eigvectors = torch.linalg.eigh(L)
        eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)
        eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
        # Retrieve eigenvalues features
        n_connected_comp, batch_eigenvalues = self.eigenvalues_features(
            eigenvalues=eigvals, num_eigenvalues=self.num_eigenvalues
        )
        # Retrieve eigenvectors features
        evector_feat = self.eigenvector_features(
            vectors=eigvectors,
            node_mask=noisy_data["node_mask"],
            n_connected=n_connected_comp,
            num_eigenvectors=self.num_eigenvectors,
        )

        evalue_feat = torch.hstack((n_connected_comp, batch_eigenvalues))
        return evalue_feat, evector_feat
//  [python]
def __call__(self, E_t, mask):
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        if self.mode == 'eigenvalues':
            eigvals = torch.symeig(L,eigenvectors=False)        # bs, n
            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp.type_as(A), batch_eigenvalues.type_as(A)

        elif self.mode == 'all':
            eigvals, eigvectors = torch.symeig(L,eigenvectors=True)
            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)
            eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
            # Retrieve eigenvalues features
            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

            # Retrieve eigenvectors features
            nonlcc_indicator, k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                                               node_mask=mask,
                                                                               n_connected=n_connected_comp)
            return n_connected_comp, batch_eigenvalues, nonlcc_indicator, k_lowest_eigenvector
        else:
            raise NotImplementedError(f"Mode {self.mode} is not implemented")
//  [python]
def __call__(self, E_t, mask):
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        if self.mode == 'eigenvalues':
            eigvals = torch.linalg.eigvalsh(L)        # bs, n
            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp.type_as(A), batch_eigenvalues.type_as(A)

        elif self.mode == 'all':
            eigvals, eigvectors = torch.linalg.eigh(L)
            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)
            eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
            # Retrieve eigenvalues features
            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

            # Retrieve eigenvectors features
            nonlcc_indicator, k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                                               node_mask=mask,
                                                                               n_connected=n_connected_comp)
            return n_connected_comp, batch_eigenvalues, nonlcc_indicator, k_lowest_eigenvector
        else:
            raise NotImplementedError(f"Mode {self.mode} is not implemented")
//  [python]
def generate_cot_eigen_vectors(args):
    print('------generate cot eigen vectors------')
    for _, subset_name in enumerate(['train', 'test']):

        subset_mesh_path = os.path.join(args.data_path, subset_name + '_norm')

        eigen_vectors_path = os.path.join(args.data_path, 'eigen_vectors', subset_name)
        eigen_values_path = os.path.join(args.data_path, 'eigen_values', subset_name)
        if not os.path.exists(eigen_vectors_path):
            os.makedirs(eigen_vectors_path)
        if not os.path.exists(eigen_values_path):
            os.makedirs(eigen_values_path)

        for file in sorted(os.listdir(subset_mesh_path)):
            file.strip()
            if os.path.splitext(file)[1] not in ['.obj']:
                continue
            print(os.path.join(subset_mesh_path, file))

            mesh = trimesh.load(os.path.join(subset_mesh_path, file), process=False)
            mesh: trimesh.Trimesh
            cot = -igl.cotmatrix(mesh.vertices, mesh.faces).toarray()
            cot = torch.from_numpy(cot).float().to(args.device)
            eigen_values, eigen_vectors = torch.symeig(cot, eigenvectors=True)
            ind = torch.argsort(eigen_values)[:]

            np.save(os.path.join(eigen_vectors_path, os.path.splitext(file)[0] + '_eigen.npy'),
                    eigen_vectors[:, ind].cpu().numpy())
            np.save(os.path.join(eigen_values_path, os.path.splitext(file)[0] + '_eigenValues.npy'),
                    eigen_values[ind].cpu().numpy())
//  [python]
def generate_cot_eigen_vectors(args):
    print('------generate cot eigen vectors------')
    for _, subset_name in enumerate(['train', 'test']):

        subset_mesh_path = os.path.join(args.data_path, subset_name + '_norm')

        eigen_vectors_path = os.path.join(args.data_path, 'eigen_vectors', subset_name)
        eigen_values_path = os.path.join(args.data_path, 'eigen_values', subset_name)
        if not os.path.exists(eigen_vectors_path):
            os.makedirs(eigen_vectors_path)
        if not os.path.exists(eigen_values_path):
            os.makedirs(eigen_values_path)

        for file in sorted(os.listdir(subset_mesh_path)):
            file.strip()
            if os.path.splitext(file)[1] not in ['.obj']:
                continue
            print(os.path.join(subset_mesh_path, file))

            mesh = trimesh.load(os.path.join(subset_mesh_path, file), process=False)
            mesh: trimesh.Trimesh
            cot = -igl.cotmatrix(mesh.vertices, mesh.faces).toarray()
            cot = torch.from_numpy(cot).float().to(args.device)
            eigen_values, eigen_vectors = torch.linalg.eigh(cot)
            ind = torch.argsort(eigen_values)[:]

            np.save(os.path.join(eigen_vectors_path, os.path.splitext(file)[0] + '_eigen.npy'),
                    eigen_vectors[:, ind].cpu().numpy())
            np.save(os.path.join(eigen_values_path, os.path.splitext(file)[0] + '_eigenValues.npy'),
                    eigen_values[ind].cpu().numpy())
//  [python]
def post_covariance_training(self, rcond: float = 1e-5, **kwargs):
        whitening_matrices = {}
        for k, cov in self.covariances.items():
            # Compute decomposition
            eigs = torch.symeig(cov)

            # Zero entries corresponding to eigenvalues smaller than rcond
            vals_rsqrt = eigs.eigenvalues.rsqrt()
            vals_rsqrt[eigs.eigenvalues < rcond * eigs.eigenvalues.max()] = 0

            # PCA whitening
            # following https://doi.org/10.1080/00031305.2016.1277159
            # and https://stats.stackexchange.com/a/594218/319192
            # but transposed (sphering with x@W instead of W@x)
            whitening_matrices[k] = eigs.eigenvectors * vals_rsqrt.unsqueeze(0)
            assert torch.allclose(
                whitening_matrices[k], eigs.eigenvectors @ vals_rsqrt.diag()
            )
        self.whitening_matrices = whitening_matrices
//  [python]
def post_covariance_training(self, rcond: float = 1e-5, **kwargs):
        whitening_matrices = {}
        for k, cov in self.covariances.items():
            # Compute decomposition
            eigs = torch.linalg.eigh(cov)

            # Zero entries corresponding to eigenvalues smaller than rcond
            vals_rsqrt = eigs.eigenvalues.rsqrt()
            vals_rsqrt[eigs.eigenvalues < rcond * eigs.eigenvalues.max()] = 0

            # PCA whitening
            # following https://doi.org/10.1080/00031305.2016.1277159
            # and https://stats.stackexchange.com/a/594218/319192
            # but transposed (sphering with x@W instead of W@x)
            whitening_matrices[k] = eigs.eigenvectors * vals_rsqrt.unsqueeze(0)
            assert torch.allclose(
                whitening_matrices[k], eigs.eigenvectors @ vals_rsqrt.diag()
            )
        self.whitening_matrices = whitening_matrices
//  [python]
def noise_estimate_batch(im, pch_size=10):
    '''
    Implement of noise level estimation of the following paper:
    Chen G , Zhu F , Heng P A . An Efficient Statistical Method for Image Noise Level Estimation[C]// 2015 IEEE International Conference
    on Computer Vision (ICCV). IEEE Computer Society, 2015.
    Input:
        im: the noise image, H x W x 3 or H x W numpy tensor, range [0,1]
        pch_size: patch_size
    Output:
        noise_level: the estimated noise level
    '''

    # if im.ndim == 3:
    #     im = im.transpose((2, 0, 1))
    # else:
    #     im = np.expand_dims(im, axis=0)

    # image to patch
    pch = im2patch_batch(im, pch_size, 3)  # C x pch_size x pch_size x num_pch tensor

    num_pch = pch.shape[4]
    pch = pch.reshape((-1, num_pch))  # d x num_pch matrix
    d = pch.shape[0]
    mu = pch.mean(dim=1, keepdim=True)  # d x 1
    X = pch - mu
    sigma_X = torch.mm(X, X.transpose(1, 0)) / num_pch
    sig_value, _ = torch.symeig(sigma_X, eigenvectors=False)
    sig_value.sort()

    for ii in range(-1, -d-1, -1):
        tau = torch.mean(sig_value[:ii])
        if torch.sum(sig_value[:ii]>tau) == torch.sum(sig_value[:ii] < tau):
            return torch.sqrt(tau)
//  [python]
def noise_estimate_batch(im, pch_size=10):
    '''
    Implement of noise level estimation of the following paper:
    Chen G , Zhu F , Heng P A . An Efficient Statistical Method for Image Noise Level Estimation[C]// 2015 IEEE International Conference
    on Computer Vision (ICCV). IEEE Computer Society, 2015.
    Input:
        im: the noise image, H x W x 3 or H x W numpy tensor, range [0,1]
        pch_size: patch_size
    Output:
        noise_level: the estimated noise level
    '''

    # if im.ndim == 3:
    #     im = im.transpose((2, 0, 1))
    # else:
    #     im = np.expand_dims(im, axis=0)

    # image to patch
    pch = im2patch_batch(im, pch_size, 3)  # C x pch_size x pch_size x num_pch tensor

    num_pch = pch.shape[4]
    pch = pch.reshape((-1, num_pch))  # d x num_pch matrix
    d = pch.shape[0]
    mu = pch.mean(dim=1, keepdim=True)  # d x 1
    X = pch - mu
    sigma_X = torch.mm(X, X.transpose(1, 0)) / num_pch
    sig_value, _ = torch.linalg.eigh(sigma_X)
    sig_value.sort()

    for ii in range(-1, -d-1, -1):
        tau = torch.mean(sig_value[:ii])
        if torch.sum(sig_value[:ii]>tau) == torch.sum(sig_value[:ii] < tau):
            return torch.sqrt(tau)
//  [python]
def smallest_eigenvector_eigh_cpu(V):
    dev = V.device

    # diagonal loading factor
    # dload = 1e-5 * torch.diag(torch.arange(V.shape[-1], device=V.device))
    # V = (V + dload).cpu()
    V = V.cpu()
    r = torch.symeig(V)

    ev = r.eigenvalues[..., 0].to(dev)
    ev = torch.reciprocal(torch.clamp(ev, min=1e-5))

    x = r.eigenvectors[..., :, 0].to(dev)

    return ev, x
//  [python]
def smallest_eigenvector_eigh_cpu(V):
    dev = V.device

    # diagonal loading factor
    # dload = 1e-5 * torch.diag(torch.arange(V.shape[-1], device=V.device))
    # V = (V + dload).cpu()
    V = V.cpu()
    r = torch.linalg.eigh(V)

    ev = r.eigenvalues[..., 0].to(dev)
    ev = torch.reciprocal(torch.clamp(ev, min=1e-5))

    x = r.eigenvectors[..., :, 0].to(dev)

    return ev, x
//  [python]
def calculate_eigenvalues_and_vectors(self):
        for layer_name, cov_matrices in self.layer_covs.items():
            eigenvalues_S, eigenvectors_S = torch.symeig(cov_matrices['S'], eigenvectors=True)
            eigenvalues_A, eigenvectors_A = torch.symeig(cov_matrices['A'], eigenvectors=True)

            self.layer_svds[layer_name] = {
                'Q_S': eigenvectors_S,
                'Q_A': eigenvectors_A,
            }
//  [python]
def calculate_eigenvalues_and_vectors(self):
        for layer_name, cov_matrices in self.layer_covs.items():
            eigenvalues_S, eigenvectors_S = torch.linalg.eigh(cov_matrices['S'], UPLO='U')
            eigenvalues_A, eigenvectors_A = torch.linalg.eigh(cov_matrices['A'], UPLO='U')

            self.layer_svds[layer_name] = {
                'Q_S': eigenvectors_S,
                'Q_A': eigenvectors_A,
            }
//  [python]
def forward(self, images, mask=None):
        B, N = images.shape[:2]
        vol_shape = images.shape[3:]

        if mask is None:
            mask = torch.ones(B, 1, *vol_shape, device=images.device, dtype=images.dtype)

        mask = mask.view(B, 1, -1)
        I = images.view(B, N, -1)
        masked_I = torch.masked_select(I, mask=mask.to(torch.bool)).view(B, N, -1)

        std, mean = torch.std_mean(masked_I, dim=-1, keepdim=True)
        X = I.sub(mean.detach()).div(std.detach())
        A = torch.bmm(X, X.transpose(1, 2)).div(X.shape[-1])

        _, Q = torch.symeig(A, eigenvectors=True)
        eigvec = Q[:, :, [-1]]

        template = I.mul(eigvec).sum(dim=1, keepdim=True).view(B, 1, *vol_shape)
        return template
//  [python]
def forward(self, images, mask=None):
        B, N = images.shape[:2]
        vol_shape = images.shape[3:]

        if mask is None:
            mask = torch.ones(B, 1, *vol_shape, device=images.device, dtype=images.dtype)

        mask = mask.view(B, 1, -1)
        I = images.view(B, N, -1)
        masked_I = torch.masked_select(I, mask=mask.to(torch.bool)).view(B, N, -1)

        std, mean = torch.std_mean(masked_I, dim=-1, keepdim=True)
        X = I.sub(mean.detach()).div(std.detach())
        A = torch.bmm(X, X.transpose(1, 2)).div(X.shape[-1])

        _, Q = torch.linalg.eigh(A)
        eigvec = Q[:, :, [-1]]

        template = I.mul(eigvec).sum(dim=1, keepdim=True).view(B, 1, *vol_shape)
        return template
//  [python]
def _compute_kfe(self, group, state):
        """Computes the covariances."""
        mod = group['mod']
        x = self.state[group['mod']]['x']
        gy = self.state[group['mod']]['gy']
        # Computation of xxt
        if group['layer_type'] == 'Conv2d':
            if not self.sua:
                x = F.conv2d(x, group['gathering_filter'],
                             stride=mod.stride, padding=mod.padding,
                             groups=mod.in_channels)
            x = x.data.permute(1, 0, 2, 3).contiguous().view(x.shape[1], -1)
        else:
            x = x.data.t()
        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones], dim=0)
        xxt = torch.mm(x, x.t()) / float(x.shape[1])
        Ex, state['kfe_x'] = torch.symeig(xxt,eigenvectors=True)
        # Computation of ggt
        if group['layer_type'] == 'Conv2d':
            gy = gy.data.permute(1, 0, 2, 3)
            state['num_locations'] = gy.shape[2] * gy.shape[3]
            gy = gy.contiguous().view(gy.shape[0], -1)
        else:
            gy = gy.data.t()
            state['num_locations'] = 1
        ggt = torch.mm(gy, gy.t()) / float(gy.shape[1])
        Eg, state['kfe_gy'] = torch.symeig(ggt,eigenvectors=True)
        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * state['num_locations']
        if group['layer_type'] == 'Conv2d' and self.sua:
            ws = group['params'][0].grad.data.size()
            state['m2'] = state['m2'].view(Eg.size(0), Ex.size(0), 1, 1).expand(-1, -1, ws[2], ws[3])
//  [python]
def _compute_kfe(self, group, state):
        """Computes the covariances."""
        mod = group['mod']
        x = self.state[group['mod']]['x']
        gy = self.state[group['mod']]['gy']
        # Computation of xxt
        if group['layer_type'] == 'Conv2d':
            if not self.sua:
                x = F.conv2d(x, group['gathering_filter'],
                             stride=mod.stride, padding=mod.padding,
                             groups=mod.in_channels)
            x = x.data.permute(1, 0, 2, 3).contiguous().view(x.shape[1], -1)
        else:
            x = x.data.t()
        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones], dim=0)
        xxt = torch.mm(x, x.t()) / float(x.shape[1])
        Ex, state['kfe_x'] = torch.linalg.eigh(xxt,UPLO='U')
        # Computation of ggt
        if group['layer_type'] == 'Conv2d':
            gy = gy.data.permute(1, 0, 2, 3)
            state['num_locations'] = gy.shape[2] * gy.shape[3]
            gy = gy.contiguous().view(gy.shape[0], -1)
        else:
            gy = gy.data.t()
            state['num_locations'] = 1
        ggt = torch.mm(gy, gy.t()) / float(gy.shape[1])
        Eg, state['kfe_gy'] = torch.linalg.eigh(ggt, UPLO='U')
        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * state['num_locations']
        if group['layer_type'] == 'Conv2d' and self.sua:
            ws = group['params'][0].grad.data.size()
            state['m2'] = state['m2'].view(Eg.size(0), Ex.size(0), 1, 1).expand(-1, -1, ws[2], ws[3])
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        """
        Returns a randomly sampled matrix on the manifold as

        .. math::

            WW^\\intercal \\qquad W_{i,j} \\sim \\texttt{init_}

        By default ``init\\_`` is a (xavier) normal distribution, so that the
        returned matrix follows a Wishart distribution.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = PSSD(layer.weight.size())
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\\_ (callable): Optional.
                    A function that takes a tensor and fills it in place according
                    to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(*(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype)
            init_(X)
            X = X @ X.transpose(-2, -1)
            L, Q = torch.symeig(X)
            L = L[..., -self.rank:]
            Q = Q[..., -self.rank:]
            if factorized:
                return L, Q
            else:
                return Q * L.unsqueeze(-2) @ Q.transpose(-2, -1)
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        """
        Returns a randomly sampled matrix on the manifold as

        .. math::

            WW^\\intercal \\qquad W_{i,j} \\sim \\texttt{init_}

        By default ``init\\_`` is a (xavier) normal distribution, so that the
        returned matrix follows a Wishart distribution.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = PSSD(layer.weight.size())
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\\_ (callable): Optional.
                    A function that takes a tensor and fills it in place according
                    to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(*(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype)
            init_(X)
            X = X @ X.transpose(-2, -1)
            L, Q = torch.linalg.eigh(X)
            L = L[..., -self.rank:]
            Q = Q[..., -self.rank:]
            if factorized:
                return L, Q
            else:
                return Q * L.unsqueeze(-2) @ Q.transpose(-2, -1)
//  [python]
def _compute_linear_pcs(self, x: torch.Tensor):
        """Computes the principal components of the data.

        Parameters
        ----------
        x
            The reference dataset.
        """
        self.means = x.mean(0)
        x = x - self.means
        cov_mat = (x.t() @ x)/(len(x)-1)
        D, V = torch.symeig(cov_mat, eigenvectors=True)
        non_zero_inds = D > self.min_eigenvalue
        self.pcs = V[:, non_zero_inds] / D[None,  non_zero_inds].sqrt()
//  [python]
def _compute_linear_pcs(self, x: torch.Tensor):
        """Computes the principal components of the data.

        Parameters
        ----------
        x
            The reference dataset.
        """
        self.means = x.mean(0)
        x = x - self.means
        cov_mat = (x.t() @ x)/(len(x)-1)
        D, V = torch.linalg.eigh(cov_mat)
        non_zero_inds = D > self.min_eigenvalue
        self.pcs = V[:, non_zero_inds] / D[None,  non_zero_inds].sqrt()
//  [python]
def _update_inv_cov(self):
        self._eigval = []
        self._eigvec = []
        for est in self._est:
            e, v = torch.symeig(est)
            self._eigval.append(e.clamp(min=0.0))
            self._eigvec.append(v)
//  [python]
def _update_inv_cov(self):
        self._eigval = []
        self._eigvec = []
        for est in self._est:
            e, v = torch.linalg.eigh(est)
            self._eigval.append(e.clamp(min=0.0))
            self._eigvec.append(v)
//  [python]
def hessian_eigen(grad: math.Tensor, hessian: math.Tensor):
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    # Transform to eigen space
    S, V = torch.symeig(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('basis'))
    V = math.reshaped_tensor(V, [*batch, *math.channel('vector,reduce')])
    V_ = math.reshaped_tensor(V, [*batch, *math.channel('reduce,vector')])
    unit = math.tensor([(1, 0), (0, 1)], math.channel(basis='x,y'), math.channel(vector=2)).vector * V_.reduce
    grad_eig = grad.vector * V.reduce  # scales like global_scale
    ig_eig = math.divide_no_nan(1., grad_eig)
    ig = ig_eig.vector * V_.reduce
    return unit, unit * S, ig
//  [python]
def hessian_eigen(grad: math.Tensor, hessian: math.Tensor):
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    # Transform to eigen space
    S, V = torch.linalg.eigh(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('basis'))
    V = math.reshaped_tensor(V, [*batch, *math.channel('vector,reduce')])
    V_ = math.reshaped_tensor(V, [*batch, *math.channel('reduce,vector')])
    unit = math.tensor([(1, 0), (0, 1)], math.channel(basis='x,y'), math.channel(vector=2)).vector * V_.reduce
    grad_eig = grad.vector * V.reduce  # scales like global_scale
    ig_eig = math.divide_no_nan(1., grad_eig)
    ig = ig_eig.vector * V_.reduce
    return unit, unit * S, ig