jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_inv_covariance(centroids, train_features, train_labels, jitters=None):
    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    cov = torch.zeros(
        centroids.shape[1], centroids.shape[1], device=centroids.device
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        for x in train_features[train_labels == c]:
            d = (x - mu_c).unsqueeze(1)
            cov += d @ d.T
    cov_scaled = cov / (train_features.shape[0] - 1)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(cov_scaled.shape[1], device=cov_scaled.device,)
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.symeig(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter
    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def compute_inv_covariance(centroids, train_features, train_labels, jitters=None):
    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    cov = torch.zeros(
        centroids.shape[1], centroids.shape[1], device=centroids.device
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        for x in train_features[train_labels == c]:
            d = (x - mu_c).unsqueeze(1)
            cov += d @ d.T
    cov_scaled = cov / (train_features.shape[0] - 1)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(cov_scaled.shape[1], device=cov_scaled.device,)
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.linalg.eigh(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter
    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def get_cone_quad(ellipse_quad, focal=-1):
    a_ = ellipse_quad[0]
    h_ = ellipse_quad[1]
    b_ = ellipse_quad[2]
    g_ = ellipse_quad[3]
    f_ = ellipse_quad[4]
    d_ = ellipse_quad[5]

    a = a_*focal**2
    b = b_*focal**2
    c = d_
    d = d_*focal**2
    f = -focal*(f_)
    g = -focal*(g_)
    h = h_*focal**2
    u = g_*focal**2
    v = f_*focal**2
    w = -focal*d_

    M = torch.tensor([[a, h, g],
                      [h, b, f],
                      [g, f, c]])

    eig_vals = torch.symeig(M)[0]

    cone_params = [a,b,c,d,f,g,h,u,v,w]
    return M, cone_params, eig_vals
//  [python]
def get_cone_quad(ellipse_quad, focal=-1):
    a_ = ellipse_quad[0]
    h_ = ellipse_quad[1]
    b_ = ellipse_quad[2]
    g_ = ellipse_quad[3]
    f_ = ellipse_quad[4]
    d_ = ellipse_quad[5]

    a = a_*focal**2
    b = b_*focal**2
    c = d_
    d = d_*focal**2
    f = -focal*(f_)
    g = -focal*(g_)
    h = h_*focal**2
    u = g_*focal**2
    v = f_*focal**2
    w = -focal*d_

    M = torch.tensor([[a, h, g],
                      [h, b, f],
                      [g, f, c]])

    eig_vals = torch.linalg.eigh(M)[0]

    cone_params = [a,b,c,d,f,g,h,u,v,w]
    return M, cone_params, eig_vals
//  [python]
def get_H(self, M):
        """
        Returns the intermediary H matrix for fitting the CCA weight
        matrices.

        M: FloatTensor (T,N)
            a centered data matrix. T is the dimension to be averaged
            over.

        Returns:
            H: FloatTensor (N,N)
        """
        MtM = (1-self.alpha)*torch.einsum("tn,tm->nm",M, M)
        rng = range(len(MtM))
        MtM[rng,rng] += self.alpha
        M = M.cpu()
        vals, vecs = torch.symeig(MtM, eigenvectors=True)
        vals[vals<self.eig_tol] = self.eig_tol
        H = torch.einsum("xt,t->xt", vecs, 1/vals.sqrt())
        H = torch.einsum("xt,wt->xw", H, vecs)
        return H
//  [python]
def get_H(self, M):
        """
        Returns the intermediary H matrix for fitting the CCA weight
        matrices.

        M: FloatTensor (T,N)
            a centered data matrix. T is the dimension to be averaged
            over.

        Returns:
            H: FloatTensor (N,N)
        """
        MtM = (1-self.alpha)*torch.einsum("tn,tm->nm",M, M)
        rng = range(len(MtM))
        MtM[rng,rng] += self.alpha
        M = M.cpu()
        vals, vecs = torch.linalg.eigh(MtM)
        vals[vals<self.eig_tol] = self.eig_tol
        H = torch.einsum("xt,t->xt", vecs, 1/vals.sqrt())
        H = torch.einsum("xt,wt->xw", H, vecs)
        return H
//  [python]
def forward(self, input):
        x = input.squeeze()
        r, d = x.size()
        G, V = torch.symeig(x, eigenvectors=True)
        tmp, order = torch.sort(G.diag(), 0, descending=True)
        G = G[order].diag()
        V = V[:,order]
        eig = [G, V]
        variables = eig
        self.save_for_backward(*variables)
        return eig
//  [python]
def forward(self, input):
        x = input.squeeze()
        r, d = x.size()
        G, V = torch.linalg.eigh(x)
        tmp, order = torch.sort(G.diag(), 0, descending=True)
        G = G[order].diag()
        V = V[:,order]
        eig = [G, V]
        variables = eig
        self.save_for_backward(*variables)
        return eig
//  [python]
def compute_multiplets(env, eps_multiplet_gap=1.0e-10):
    D= torch.zeros(env.chi+1, dtype=env.dtype, device=env.device)
    D[:env.chi], U= torch.symeig(env.C[env.keyC])
    D, p= torch.sort(torch.abs(D),descending=True)
    m=[]
    l=0
    for i in range(env.chi):
        l+=1
        g=D[i]-D[i+1]
        #print(f"{i} {D[i]} {g}", end=" ")
        if g>eps_multiplet_gap:
            #print(f"{l}", end=" ")
            m.append(l)
            l=0
    return m
//  [python]
def compute_multiplets(env, eps_multiplet_gap=1.0e-10):
    D= torch.zeros(env.chi+1, dtype=env.dtype, device=env.device)
    D[:env.chi], U= torch.linalg.eigh(env.C[env.keyC])
    D, p= torch.sort(torch.abs(D),descending=True)
    m=[]
    l=0
    for i in range(env.chi):
        l+=1
        g=D[i]-D[i+1]
        #print(f"{i} {D[i]} {g}", end=" ")
        if g>eps_multiplet_gap:
            #print(f"{l}", end=" ")
            m.append(l)
            l=0
    return m
//  [python]
def forward_torch(self, target):
        """
            Full-diagonalization of Hamiltonian using the "symeig" function
        built in Pytorch.
        """
        V = torch.diag(self.potential)
        H = self.K + V
        _, psis = torch.symeig(H, eigenvectors=True)
        self.psi0 = psis[:, 0]
        loss = 1. - (self.psi0.abs() * target).sum()
        return loss
//  [python]
def forward_torch(self, target):
        """
            Full-diagonalization of Hamiltonian using the "symeig" function
        built in Pytorch.
        """
        V = torch.diag(self.potential)
        H = self.K + V
        _, psis = torch.linalg.eigh(H)
        self.psi0 = psis[:, 0]
        loss = 1. - (self.psi0.abs() * target).sum()
        return loss
//  [python]
def get_wuw_mats(self, c00):
        """Calculate whitening matrix and unwhitening matrix.
           Method adapted from deeptime (https://github.com/markovmodel/deeptime/blob/master/time-lagged-autoencoder/tae/utils.py)

        Parameters
        ----------
        c00 : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            Covariance matrix

        Returns
        -------
        uwm : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            unwhitening matrix
        wm : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            whitening matrix
        """
        # Previous implementation
        # uwm = sqrtm(c00).real
        # wm = inv(uwm).real
        # return uwm, wm

        # Updated implementation
        e, v = torch.symeig(torch.from_numpy(c00).double(), eigenvectors=True)
        # In valid covariance matrix the smallest eigenvalue should be positive
        # because the covariance matrix is a positive semidefinite matrix
        # https://stats.stackexchange.com/questions/52976/is-a-sample-covariance-matrix-always-symmetric-and-positive-definite
        assert torch.min(e) > 0.0
        d = torch.diag(1.0 / torch.sqrt(e))
        wm = torch.mm(torch.mm(v, d), v.t())
        return inv(wm.numpy()), wm.numpy()
//  [python]
def get_wuw_mats(self, c00):
        """Calculate whitening matrix and unwhitening matrix.
           Method adapted from deeptime (https://github.com/markovmodel/deeptime/blob/master/time-lagged-autoencoder/tae/utils.py)

        Parameters
        ----------
        c00 : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            Covariance matrix

        Returns
        -------
        uwm : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            unwhitening matrix
        wm : np.ndarray, shape=(n_atoms*3,n_atoms*3)
            whitening matrix
        """
        # Previous implementation
        # uwm = sqrtm(c00).real
        # wm = inv(uwm).real
        # return uwm, wm

        # Updated implementation
        e, v = torch.linalg.eigh(torch.from_numpy(c00).double())
        # In valid covariance matrix the smallest eigenvalue should be positive
        # because the covariance matrix is a positive semidefinite matrix
        # https://stats.stackexchange.com/questions/52976/is-a-sample-covariance-matrix-always-symmetric-and-positive-definite
        assert torch.min(e) > 0.0
        d = torch.diag(1.0 / torch.sqrt(e))
        wm = torch.mm(torch.mm(v, d), v.t())
        return inv(wm.numpy()), wm.numpy()
//  [python]
def zca_from_cov(cov):
    evals, evecs = torch.symeig(cov.double(), eigenvectors=True)
    zca = torch.mm(torch.mm(evecs, torch.diag
                            (evals.sqrt().clamp(1e-20).reciprocal())),
                   evecs.t()).to(cov.dtype)
    return zca
//  [python]
def zca_from_cov(cov):
    evals, evecs = torch.linalg.eigh(cov.double())
    zca = torch.mm(torch.mm(evecs, torch.diag
                            (evals.sqrt().clamp(1e-20).reciprocal())),
                   evecs.t()).to(cov.dtype)
    return zca
//  [python]
def estimate_normals(
    q_points: Tensor, s_points: Optional[Tensor] = None, k: int = 50, disambiguate_directions: bool = True
):
    """Estimate normals for the query points from the support points.

    Modified from [PyTorch3d](https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/ops/points_normals.py).

    Args:
        q_points (Tensor): query points (N, 3)
        s_points (Tensor=None): support points. If None, use q_points (M, 3)
        k (int=50): number of neighbors
        disambiguate_directions (bool=True): disambiguate directions.

    Returns:
        q_normals (Tensor): normals of the query points (N, 3)
    """
    if s_points is None:
        s_points = q_points

    knn_indices = knn(q_points, s_points, k)  # (N, K)
    knn_points = index_select(s_points, knn_indices, dim=0)  # (N, K, 3)
    centroids = knn_points.mean(dim=1, keepdim=True)  # (N, 1, 3)
    knn_offsets = knn_points - centroids  # (N, K, 3)

    cov_mat = knn_offsets.unsqueeze(3) * knn_offsets.unsqueeze(2)  # (N, K, 3, 1) x (N, K, 1, 3)
    cov_mat = cov_mat.mean(dim=1)  # (N, 3, 3)
    _, eigenvectors = torch.symeig(cov_mat.cpu(), eigenvectors=True)  # (N, 3), (N, 3, 3)
    eigenvectors = eigenvectors.cuda()
    normals = eigenvectors[:, :, 0]  # (N, 3)

    if disambiguate_directions:
        knn_offsets = knn_points - q_points.unsqueeze(1)  # (N, K, 3) - (N, 3) -> (N, 1, 3)
        projections = (normals.unsqueeze(1) * knn_offsets).sum(2)  # (N, 3) -> (N, 1, 3) x (N, K, 3) -> (N, K)
        flips = torch.gt(projections.sum(dim=-1, keepdim=True), 0).float()  # (N, K) -> (N, 1)
        # flips = torch.lt(torch.gt(projections, 0).float().mean(dim=1, keepdim=True), 0.5).float()  # (N, K) -> (N, 1)
        normals = (1.0 - 2.0 * flips) * normals

    return normals
//  [python]
def estimate_normals(
    q_points: Tensor, s_points: Optional[Tensor] = None, k: int = 50, disambiguate_directions: bool = True
):
    """Estimate normals for the query points from the support points.

    Modified from [PyTorch3d](https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/ops/points_normals.py).

    Args:
        q_points (Tensor): query points (N, 3)
        s_points (Tensor=None): support points. If None, use q_points (M, 3)
        k (int=50): number of neighbors
        disambiguate_directions (bool=True): disambiguate directions.

    Returns:
        q_normals (Tensor): normals of the query points (N, 3)
    """
    if s_points is None:
        s_points = q_points

    knn_indices = knn(q_points, s_points, k)  # (N, K)
    knn_points = index_select(s_points, knn_indices, dim=0)  # (N, K, 3)
    centroids = knn_points.mean(dim=1, keepdim=True)  # (N, 1, 3)
    knn_offsets = knn_points - centroids  # (N, K, 3)

    cov_mat = knn_offsets.unsqueeze(3) * knn_offsets.unsqueeze(2)  # (N, K, 3, 1) x (N, K, 1, 3)
    cov_mat = cov_mat.mean(dim=1)  # (N, 3, 3)
    eigenvalues, eigenvectors = torch.linalg.eigh(cov_mat.cpu())  # (N, 3), (N, 3, 3)
    eigenvectors = eigenvectors.cuda()
    normals = eigenvectors[:, :, 0]  # (N, 3)

    if disambiguate_directions:
        knn_offsets = knn_points - q_points.unsqueeze(1)  # (N, K, 3) - (N, 3) -> (N, 1, 3)
        projections = (normals.unsqueeze(1) * knn_offsets).sum(2)  # (N, 3) -> (N, 1, 3) x (N, K, 3) -> (N, K)
        flips = torch.gt(projections.sum(dim=-1, keepdim=True), 0).float()  # (N, K) -> (N, 1)
        # flips = torch.lt(torch.gt(projections, 0).float().mean(dim=1, keepdim=True), 0.5).float()  # (N, K) -> (N, 1)
        normals = (1.0 - 2.0 * flips) * normals

    return normals
//  [python]
def loss(self, H1, H2):
        """
        Compute the loss (negative correlation) between 2 views. Details can
        be found in [#1DCCA]_ or the documentation for ``DCCA``.

        Parameters
        ----------
        H1: torch.tensor, shape (n_samples, n_features)
            View 1 data.
        H2: torch.tensor, shape (n_samples, n_features)
            View 2 data.
        """

        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        # Transpose matrices so each column is a sample
        H1, H2 = H1.t(), H2.t()

        o1 = o2 = H1.size(0)

        m = H1.size(1)

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

        # Compute covariance matrices and add diagonal so they are
        # positive definite
        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=self.device_)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + \
            r2 * torch.eye(o2, device=self.device_)

        # Calculate the root inverse of covariance matrices by using
        # eigen decomposition
        [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
        [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

        # Additional code to increase numerical stability
        posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
        D1 = D1[posInd1]
        V1 = V1[:, posInd1]
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        D2 = D2[posInd2]
        V2 = V2[:, posInd2]

        # Compute sigma hat matrices using the edited covariance matrices
        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())

        # Compute the T matrix, whose matrix trace norm is the loss
        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)

        if self.use_all_singular_values_:
            # all singular values are used to calculate the correlation (and
            # thus the loss as well)
            tmp = torch.trace(torch.matmul(Tval.t(), Tval))
            corr = torch.sqrt(tmp)
        else:
            # just the top self.n_components_ singular values are used to
            # compute the loss
            U, V = torch.symeig(torch.matmul(
                Tval.t(), Tval), eigenvectors=True)
            U = U.topk(self.n_components_)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def loss(self, H1, H2):
        """
        Compute the loss (negative correlation) between 2 views. Details can
        be found in [#1DCCA]_ or the documentation for ``DCCA``.

        Parameters
        ----------
        H1: torch.tensor, shape (n_samples, n_features)
            View 1 data.
        H2: torch.tensor, shape (n_samples, n_features)
            View 2 data.
        """

        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        # Transpose matrices so each column is a sample
        H1, H2 = H1.t(), H2.t()

        o1 = o2 = H1.size(0)

        m = H1.size(1)

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

        # Compute covariance matrices and add diagonal so they are
        # positive definite
        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=self.device_)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + \
            r2 * torch.eye(o2, device=self.device_)

        # Calculate the root inverse of covariance matrices by using
        # eigen decomposition
        [D1, V1] = torch.linalg.eigh(SigmaHat11)
        [D2, V2] = torch.linalg.eigh(SigmaHat22)

        # Additional code to increase numerical stability
        posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
        D1 = D1[posInd1]
        V1 = V1[:, posInd1]
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        D2 = D2[posInd2]
        V2 = V2[:, posInd2]

        # Compute sigma hat matrices using the edited covariance matrices
        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())

        # Compute the T matrix, whose matrix trace norm is the loss
        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)

        if self.use_all_singular_values_:
            # all singular values are used to calculate the correlation (and
            # thus the loss as well)
            tmp = torch.trace(torch.matmul(Tval.t(), Tval))
            corr = torch.sqrt(tmp)
        else:
            # just the top self.n_components_ singular values are used to
            # compute the loss
            U, V = torch.linalg.eigh(torch.matmul(
                Tval.t(), Tval))
            U = U.topk(self.n_components_)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def test_symeig(self):
        # Small case
        tensor = torch.randn(3, 3).cuda()
        tensor = torch.mm(tensor, tensor.t())
        eigval, eigvec = torch.symeig(tensor, eigenvectors=True)
        self.assertEqual(tensor, torch.mm(torch.mm(eigvec, eigval.diag()), eigvec.t()))

        # Large case
        tensor = torch.randn(257, 257).cuda()
        tensor = torch.mm(tensor, tensor.t())
        eigval, eigvec = torch.symeig(tensor, eigenvectors=True)
        self.assertEqual(tensor, torch.mm(torch.mm(eigvec, eigval.diag()), eigvec.t()))
//  [python]
def test_symeig(self):
        # Small case
        tensor = torch.randn(3, 3).cuda()
        tensor = torch.mm(tensor, tensor.t())
        eigval, eigvec = torch.linalg.eigh(tensor)
        self.assertEqual(tensor, torch.mm(torch.mm(eigvec, eigval.diag()), eigvec.t()))

        # Large case
        tensor = torch.randn(257, 257).cuda()
        tensor = torch.mm(tensor, tensor.t())
        eigval, eigvec = torch.linalg.eigh(tensor)
        self.assertEqual(tensor, torch.mm(torch.mm(eigvec, eigval.diag()), eigvec.t()))
//  [python]
def _max_abs_eig(self, A):
        l = torch.symeig(A)[0]
        return torch.max(torch.abs(l))
//  [python]
def _max_abs_eig(self, A):
        l = torch.linalg.eigh(A)[0]
        return torch.max(torch.abs(l))
//  [python]
def _update_decomps(self, module, n):
        eps = 1e-10  # for numerical stability
        D_inp, Q_inp = torch.symeig(self.inputs_outer[module][n], eigenvectors=True)
        try:
            D_grad, Q_grad = torch.symeig(self.grad_ouputs_outer[module][n], eigenvectors=True)
        except Exception as e:
            print(self.grad_ouputs_outer[module][n].min(),
                  self.grad_ouputs_outer[module][n].abs().min(),
                  self.grad_ouputs_outer[module][n].max())
            raise e

        D_inp.mul_((D_inp > eps).float())
        D_grad.mul_((D_grad > eps).float())

        self.D_inputs[module][n], self.Q_inputs[module][n] = D_inp, Q_inp
        self.D_grads[module][n], self.Q_grads[module][n] = D_grad, Q_grad
//  [python]
def _update_decomps(self, module, n):
        eps = 1e-10  # for numerical stability
        D_inp, Q_inp = torch.linalg.eigh(self.inputs_outer[module][n])
        try:
            D_grad, Q_grad = torch.linalg.eigh(self.grad_ouputs_outer[module][n])
        except Exception as e:
            print(self.grad_ouputs_outer[module][n].min(),
                  self.grad_ouputs_outer[module][n].abs().min(),
                  self.grad_ouputs_outer[module][n].max())
            raise e

        D_inp.mul_((D_inp > eps).float())
        D_grad.mul_((D_grad > eps).float())

        self.D_inputs[module][n], self.Q_inputs[module][n] = D_inp, Q_inp
        self.D_grads[module][n], self.Q_grads[module][n] = D_grad, Q_grad
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor([[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0], [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
                          [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.symeig(K, True, False)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor([[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0], [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
                          [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.linalg.eigh(K)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def _apply_floor_scalar(self, A, b):
        l, U = torch.symeig(A, eigenvectors=True)
        num_floored = torch.sum(l < b).item()
        l = torch.clamp(l, min=b)
        A = torch.matmul(U, l.unsqueeze(1) * U.t())
        return A, num_floored
//  [python]
def _apply_floor_scalar(self, A, b):
        l, U = torch.linalg.eigh(A)
        num_floored = torch.sum(l < b).item()
        l = torch.clamp(l, min=b)
        A = torch.matmul(U, l.unsqueeze(1) * U.t())
        return A, num_floored
//  [python]
def _update_inv(self, m):
        """Do eigen decomposition for computing inverse of the ~ fisher.
        :param m: The layer
        :return: no returns.
        """
        eps = 1e-10  # for numerical stability
        self.d_a[m], self.Q_a[m] = torch.symeig(
            self.m_aa[m], eigenvectors=True)
        self.d_g[m], self.Q_g[m] = torch.symeig(
            self.m_gg[m], eigenvectors=True)

        self.d_a[m].mul_((self.d_a[m] > eps).float())
        self.d_g[m].mul_((self.d_g[m] > eps).float())
//  [python]
def _update_inv(self, m):
        """Do eigen decomposition for computing inverse of the ~ fisher.
        :param m: The layer
        :return: no returns.
        """
        eps = 1e-10  # for numerical stability
        self.d_a[m], self.Q_a[m] = torch.linalg.eigh(
            self.m_aa[m])
        self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
            self.m_gg[m])

        self.d_a[m].mul_((self.d_a[m] > eps).float())
        self.d_g[m].mul_((self.d_g[m] > eps).float())
//  [python]
E, Z = torch.symeig(A, eigenvectors, True)
E, Z = xxx(A, eigenvectors, True)
//  [python]
E, Z = torch.linalg.eigh(A, eigenvectors=True)
//  [python]
def _minimum_divergence_whitening(self, h, H, component_batches):
        print('Minimum divergence re-estimation...')
        l, U = torch.symeig(H, eigenvectors=True)
        l = torch.clamp(l, min=1e-7)
        P1 = torch.rsqrt(l) * U  # transposed
        torch.matmul(h, P1, out=h)  # In place operation, so that the result is available for update_means()
        if self.prior_offset != 0:  # Augmented formulation
            self.prior_offset = h[0]
            print('Prior offset: {}'.format(self.prior_offset))
        P1 = torch.inverse(P1)
        #P1 = torch.pinverse(P1, 0.0001)
        for bstart, bend in component_batches:
            self.t_matrix[bstart:bend, :, :] = P1.matmul(self.t_matrix[bstart:bend, :, :])
//  [python]
def _minimum_divergence_whitening(self, h, H, component_batches):
        print('Minimum divergence re-estimation...')
        l, U = torch.linalg.eigh(H)
        l = torch.clamp(l, min=1e-7)
        P1 = torch.rsqrt(l) * U  # transposed
        torch.matmul(h, P1, out=h)  # In place operation, so that the result is available for update_means()
        if self.prior_offset != 0:  # Augmented formulation
            self.prior_offset = h[0]
            print('Prior offset: {}'.format(self.prior_offset))
        P1 = torch.inverse(P1)
        #P1 = torch.pinverse(P1, 0.0001)
        for bstart, bend in component_batches:
            self.t_matrix[bstart:bend, :, :] = P1.matmul(self.t_matrix[bstart:bend, :, :])
//  [python]
def sync_perm(lap: torch.Tensor, d: int, t: float = -np.inf):
    e, V = torch.symeig(lap, eigenvectors=True)
    V = V[..., :d]
    Zd = V.bmm(V.transpose(-1, -2))
    if t > 1e-5:
        Zd[Zd < t] = 0.0
    return Zd
//  [python]
def sync_perm(lap: torch.Tensor, d: int, t: float = -np.inf):
    e, V = torch.linalg.eigh(lap)
    V = V[..., :d]
    Zd = V.bmm(V.transpose(-1, -2))
    if t > 1e-5:
        Zd[Zd < t] = 0.0
    return Zd
//  [python]
def avg_quaternions_eigen(quats: torch.Tensor, weights=None) -> torch.Tensor:
    # Why flipped?
    raise NotImplementedError
    if weights is not None and len(quats) != len(weights):
        raise ValueError("Args are of different length")
    if weights is None:
        weights = torch.ones_like(quats[:, 0])
    accum = torch.zeros((4, 4), device=quats.device)
    for i, q in enumerate(quats):
        qOuterProdWeighted = torch.outer(q, q) * weights[i]
        accum += qOuterProdWeighted
    _, eigVecs = torch.symeig(accum, eigenvectors=True)
    return eigVecs[:, 0]
//  [python]
def avg_quaternions_eigen(quats: torch.Tensor, weights=None) -> torch.Tensor:
    # Why flipped?
    raise NotImplementedError
    if weights is not None and len(quats) != len(weights):
        raise ValueError("Args are of different length")
    if weights is None:
        weights = torch.ones_like(quats[:, 0])
    accum = torch.zeros((4, 4), device=quats.device)
    for i, q in enumerate(quats):
        qOuterProdWeighted = torch.outer(q, q) * weights[i]
        accum += qOuterProdWeighted
    _, eigVecs = torch.linalg.eigh(accum)
    return eigVecs[:, 0]
//  [python]
def compute_inv_covariance(centroids, train_features, train_labels, jitters=None):
    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    cov = torch.zeros(
        centroids.shape[1], centroids.shape[1], device=centroids.device
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        for x in train_features[train_labels == c]:
            d = (x - mu_c).unsqueeze(1)
            cov += d @ d.T
    cov_scaled = cov / (train_features.shape[0] - 1)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(
            cov_scaled.shape[1],
            device=cov_scaled.device,
        )
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.symeig(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter
    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def compute_inv_covariance(centroids, train_features, train_labels, jitters=None):
    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    cov = torch.zeros(
        centroids.shape[1], centroids.shape[1], device=centroids.device
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        for x in train_features[train_labels == c]:
            d = (x - mu_c).unsqueeze(1)
            cov += d @ d.T
    cov_scaled = cov / (train_features.shape[0] - 1)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(
            cov_scaled.shape[1],
            device=cov_scaled.device,
        )
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.linalg.eigh(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter
    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, 'Can handle only one parameter at the moment'
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.symeig(self.m_aa[m],
                                                        eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.symeig(self.m_gg[m],
                                                        eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) +
                       la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, 'Can handle only one parameter at the moment'
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(self.m_aa[m],
                                                        eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(self.m_gg[m],
                                                        eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) +
                       la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def forward(self, A, ad_decomp_reg):
        r"""
        :param A: square symmetric matrix
        :type A: torch.tensor
        :return: eigenvalues values D, eigenvectors vectors U
        :rtype: torch.tensor, torch.tensor

        Computes symmetric decomposition :math:`M= UDU^T`.
        """
        # input validation (A is square and symmetric) is provided by torch.symeig
        
        D, U = torch.symeig(A, eigenvectors=True)
        # torch.symeig returns eigenpairs ordered in the ascending order with 
        # respect to eigenvalues. Reorder the eigenpairs by abs value of the eigenvalues
        # abs(D)
        absD,p= torch.sort(torch.abs(D),descending=True)
        D= D[p]
        U= U[:,p]
        
        self.save_for_backward(D,U,ad_decomp_reg)
        return D,U
//  [python]
def forward(self, A, ad_decomp_reg):
        r"""
        :param A: square symmetric matrix
        :type A: torch.tensor
        :return: eigenvalues values D, eigenvectors vectors U
        :rtype: torch.tensor, torch.tensor

        Computes symmetric decomposition :math:`M= UDU^T`.
        """
        # input validation (A is square and symmetric) is provided by torch.linalg.eigh
        
        D, U = torch.linalg.eigh(A)
        # torch.linalg.eigh returns eigenpairs ordered in the ascending order with 
        # respect to eigenvalues. Reorder the eigenpairs by abs value of the eigenvalues
        # abs(D)
        absD,p= torch.sort(torch.abs(D),descending=True)
        D= D[p]
        U= U[:,p]
        
        self.save_for_backward(D,U,ad_decomp_reg)
        return D,U
//  [python]
def sync_motion_seg(z_mat: torch.Tensor, force_d: int = -1, t: float = -np.inf, cut_thres: float = 0.1):
    e, V = torch.symeig(z_mat, eigenvectors=True)
    if force_d != -1:
        d = force_d
    else:
        assert z_mat.size(0) == 1
        e_leading = e[:, -10:]
        total_est_points = torch.sum(e_leading, -1)
        e_th = total_est_points * cut_thres
        e_count = torch.sum(e.detach() > e_th, dim=1)
        d = e_count.max().item()

    V = V[..., -d:]  # (B, M, d)
    e = e[..., -d:]  # (B, d)
    V = V * e.sqrt().unsqueeze(1)
    v_sign = V.detach().sum(dim=1, keepdim=True).sign()
    V = V * v_sign
    if t > -1e5:
        V[V < t] = 0.0

    return V
//  [python]
def sync_motion_seg(z_mat: torch.Tensor, force_d: int = -1, t: float = -np.inf, cut_thres: float = 0.1):
    e, V = torch.linalg.eigh(z_mat)
    if force_d != -1:
        d = force_d
    else:
        assert z_mat.size(0) == 1
        e_leading = e[:, -10:]
        total_est_points = torch.sum(e_leading, -1)
        e_th = total_est_points * cut_thres
        e_count = torch.sum(e.detach() > e_th, dim=1)
        d = e_count.max().item()

    V = V[..., -d:]  # (B, M, d)
    e = e[..., -d:]  # (B, d)
    V = V * e.sqrt().unsqueeze(1)
    v_sign = V.detach().sum(dim=1, keepdim=True).sign()
    V = V * v_sign
    if t > -1e5:
        V[V < t] = 0.0

    return V
//  [python]
def get_grad(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)
//  [python]
def get_grad(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(
                    self.m_aa[m])
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
                    self.m_gg[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)
//  [python]
def test_SYMEIG_rank_deficient():
    m= 50
    r= 10
    M= torch.rand((m,m),dtype=torch.float64)
    M= M+M.t()
    D,U= torch.symeig(M, eigenvectors=True)
    D[-r:]=0
    M= U@torch.diag(D)@U.t()

    D,U= SYMEIG.apply(M)
    assert( torch.norm(M-U@torch.diag(D)@U.t()) < D[0]*(M.size()[0]**2)*1e-14 )

    M.requires_grad_(True)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        D,U= SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_rank_deficient():
    m= 50
    r= 10
    M= torch.rand((m,m),dtype=torch.float64)
    M= M+M.t()
    D,U= torch.linalg.eigh(M)
    D[-r:]=0
    M= U@torch.diag(D)@U.t()

    D,U= SYMEIG.apply(M)
    assert( torch.norm(M-U@torch.diag(D)@U.t()) < D[0]*(M.size()[0]**2)*1e-14 )

    M.requires_grad_(True)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        D,U= SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def get_eigenvalues(factors: List[Tensor],
                    verbose: bool = False) -> Tensor:
    """Computes the eigenvalues of KFAC, EFB or diagonal factors.

    Args:
        factors: A list of KFAC, EFB or diagonal factors.
        verbose: Prints out progress if True.

    Returns:
        The eigenvalues of all KFAC, EFB or diagonal factors.
    """
    eigenvalues = Tensor()
    factors = tqdm(factors, disable=not verbose)
    for layer, factor in enumerate(factors):
        factors.set_description(desc=f"Layer [{layer + 1}/{len(factors)}]")
        if len(factor) == 2:
            xxt_eigvals = torch.symeig(factor[0])[0]
            ggt_eigvals = torch.symeig(factor[1])[0]
            eigenvalues = torch.cat([eigenvalues, torch.ger(xxt_eigvals, ggt_eigvals).contiguous().view(-1)])
        else:
            eigenvalues = torch.cat([eigenvalues, factor.contiguous().view(-1)])
    return eigenvalues
//  [python]
def get_eigenvalues(factors: List[Tensor],
                    verbose: bool = False) -> Tensor:
    """Computes the eigenvalues of KFAC, EFB or diagonal factors.

    Args:
        factors: A list of KFAC, EFB or diagonal factors.
        verbose: Prints out progress if True.

    Returns:
        The eigenvalues of all KFAC, EFB or diagonal factors.
    """
    eigenvalues = Tensor()
    factors = tqdm(factors, disable=not verbose)
    for layer, factor in enumerate(factors):
        factors.set_description(desc=f"Layer [{layer + 1}/{len(factors)}]")
        if len(factor) == 2:
            xxt_eigvals = torch.linalg.eigh(factor[0])[0]
            ggt_eigvals = torch.linalg.eigh(factor[1])[0]
            eigenvalues = torch.cat([eigenvalues, torch.ger(xxt_eigvals, ggt_eigvals).contiguous().view(-1)])
        else:
            eigenvalues = torch.cat([eigenvalues, factor.contiguous().view(-1)])
    return eigenvalues
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                print('##############', self.steps)
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                print(self.m_aa.keys())
                for k in self.m_aa.keys():
                    print(k == m)
                xxxx, yyyy = torch.symeig(
                    self.m_aa[m], eigenvectors=True)
                self.d_a[m], self.Q_a[m] = xxxx, yyyy
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        # self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                print('##############', self.steps)
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                print(self.m_aa.keys())
                for k in self.m_aa.keys():
                    print(k == m)
                xxxx, yyyy = torch.linalg.eigh(
                    self.m_aa[m], eigenvectors=True)
                self.d_a[m], self.Q_a[m] = xxxx, yyyy
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        # self.optim.step()
        self.steps += 1