jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_eig(self, device, dtype):
        a = torch.Tensor(((1.96, 0.00, 0.00, 0.00, 0.00),
                          (-6.49, 3.80, 0.00, 0.00, 0.00),
                          (-0.47, -6.39, 4.17, 0.00, 0.00),
                          (-7.20, 1.50, -1.51, 5.70, 0.00),
                          (-0.65, -6.34, 2.67, 1.80, -7.10))).t().contiguous().to(dtype=dtype, device=device)
        e = torch.eig(a)[0]
        ee, vv = torch.eig(a, True)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)

        # test reuse
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        # test non-contiguous
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')

        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'A should be 2 dimensional',
            lambda: torch.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should be square',
            lambda: torch.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should not contain infs or NaNs',
            lambda: torch.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should not contain infs or NaNs',
            lambda: torch.eig(np.nan * torch.ones((2, 2))))
//  [python]
def test_eig(self, device, dtype):
        a = torch.Tensor(((1.96, 0.00, 0.00, 0.00, 0.00),
                          (-6.49, 3.80, 0.00, 0.00, 0.00),
                          (-0.47, -6.39, 4.17, 0.00, 0.00),
                          (-7.20, 1.50, -1.51, 5.70, 0.00),
                          (-0.65, -6.34, 2.67, 1.80, -7.10))).t().contiguous().to(dtype=dtype, device=device)
        e = torch.linalg.eig(a)[0]
        ee, vv = torch.linalg.eig(a)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.linalg.eig(a, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)

        # test reuse
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.linalg.eig(X, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        torch.linalg.eig(X, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        # test non-contiguous
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.linalg.eig(X, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, atol=1e-8, rtol=0, msg='VeV\' wrong')

        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'A should be 2 dimensional',
            lambda: torch.linalg.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should be square',
            lambda: torch.linalg.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'A should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.nan * torch.ones((2, 2))))
//  [python]
def compute_eigenvalues(covar):
    '''
    :param covar: 2 x 2
    :return: eigenvalues, eigenvectors, orientation
    '''
    # eigenvalues e: 2 x 2, each row is an eigenvalue e[i,0]+j*e[i, 1],
    # eigenvectors v: 2 x 2, each column is a corresponding eigenvector v[:, i]
    e, v = torch.eig(covar, eigenvectors=True)
    _, indices = torch.sort(e[:, 0], descending=True, dim=0)
    e2, v2 = e[indices, :], v[:, indices]
    radian = torch.atan2(v2[1, 0], v2[0, 0])  # atan2(vy, vx)
    angle = radian / 3.1415926 * 180  # orientation of the eigenvector for major axis
    return e2, v2, angle
//  [python]
def compute_eigenvalues(covar):
    '''
    :param covar: 2 x 2
    :return: eigenvalues, eigenvectors, orientation
    '''
    # eigenvalues e: 2 x 2, each row is an eigenvalue e[i,0]+j*e[i, 1],
    # eigenvectors v: 2 x 2, each column is a corresponding eigenvector v[:, i]
    e, v = torch.linalg.eig(covar)
    _, indices = torch.sort(e[:, 0], descending=True, dim=0)
    e2, v2 = e[indices, :], v[:, indices]
    radian = torch.atan2(v2[1, 0], v2[0, 0])  # atan2(vy, vx)
    angle = radian / 3.1415926 * 180  # orientation of the eigenvector for major axis
    return e2, v2, angle
//  [python]
def wasserstein2_loss(mean_reals: torch.Tensor,
                      mean_gen: torch.Tensor,
                      cov_reals: torch.Tensor,
                      cov_gen: torch.Tensor,
                      eps: float = 1e-12) -> torch.Tensor:
    """Computes 2-Wasserstein distance."""
    mean_term = torch.sum(torch.square(mean_reals - mean_gen.squeeze(0)))
    eigenvalues, _ = torch.eig(torch.matmul(cov_gen, cov_reals), eigenvectors=True)  # Eigenvalues shape: (D, 2) (real and imaginary parts).
    cov_term = torch.trace(cov_reals) + torch.trace(cov_gen) - 2 * torch.sum(torch.sqrt(eigenvalues[:, 0] + eps))
    return mean_term + cov_term
//  [python]
def wasserstein2_loss(mean_reals: torch.Tensor,
                      mean_gen: torch.Tensor,
                      cov_reals: torch.Tensor,
                      cov_gen: torch.Tensor,
                      eps: float = 1e-12) -> torch.Tensor:
    """Computes 2-Wasserstein distance."""
    mean_term = torch.sum(torch.square(mean_reals - mean_gen.squeeze(0)))
    eigenvalues = torch.linalg.eigvals(torch.matmul(cov_gen, cov_reals))  # Eigenvalues shape: (D,)
    cov_term = torch.trace(cov_reals) + torch.trace(cov_gen) - 2 * torch.sum(torch.sqrt(eigenvalues.real + eps))
    return mean_term + cov_term
//  [python]
def eig_real(mat: torch.Tensor) -> torch.Tensor:
    """Wrapper around torch.eig which discards imaginary values and returns result in descending order.

    Prints warning when non-zero imaginary parts detected, see "Criteria for the reality of matrix eigenvalues"
    Products of symmetric matrices are not symmetric but have eigenvalues for no imaginary parts.
    https://link.springer.com/article/10.1007%2FBF01195188
    """

    evals = torch.eig(mat).eigenvalues
    re_part = evals[:, 0]  # extract real part
    im_part = evals[:, 1]  # extract real part
    if im_part.sum() / evals.max() > 1e-7:
        print("Warning, eig_real is discarding non-zero imaginary parts")
    re_part = re_part.sort(descending=True).values
    return re_part
//  [python]
def eig_real(mat: torch.Tensor) -> torch.Tensor:
    """Wrapper around torch.eig which discards imaginary values and returns result in descending order.

    Prints warning when non-zero imaginary parts detected, see "Criteria for the reality of matrix eigenvalues"
    Products of symmetric matrices are not symmetric but have eigenvalues for no imaginary parts.
    https://link.springer.com/article/10.1007%2FBF01195188
    """

    evals = torch.linalg.eig(mat).eigenvalues
    re_part = evals[:, 0]  # extract real part
    im_part = evals[:, 1]  # extract real part
    if im_part.sum() / evals.max() > 1e-7:
        print("Warning, eig_real is discarding non-zero imaginary parts")
    re_part = re_part.sort(descending=True).values
    return re_part
//  [python]
def roots(p):
    n = p.size(-1)
    A = torch.diag(torch.ones(n - 2), -1)
    A[0, :] = -p[1:] / p[0]
    # r = torch.symeig(A)
    r = torch.eig(A)

    r_eig = np.linalg.eigvals(A)
    r_np = np.roots(p.cpu().numpy())

    return r
//  [python]
def roots(p):
    n = p.size(-1)
    A = torch.diag(torch.ones(n - 2), -1)
    A[0, :] = -p[1:] / p[0]
    # r = torch.symeig(A)
    r = torch.linalg.eig(A)

    r_eig = np.linalg.eigvals(A)
    r_np = np.roots(p.cpu().numpy())

    return r
//  [python]
def PCA(data, k=2):
    X = torch.from_numpy(data)
    X_mean = torch.mean(X, 0)
    X = X - X_mean.expand_as(X)

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

    v, w = torch.linalg.eig(torch.mm(X, torch.t(X)))
    return torch.mm(w[:, :k], X)
//  [python]
def __init__(self, dim, mean=None, cov=None):
        super().__init__(dim)
        self._has_mean = mean is not None
        if self._has_mean:
            assert len(mean.shape) == 1,\
                "`mean` must be a vector"
            assert mean.shape[-1] == self.dim, "`mean` must have dimension `dim`"
            self.register_buffer("_mean", mean.unsqueeze(0))
        self._has_cov = cov is not None
        if self._has_cov:
            assert len(cov.shape) == 2 and cov.shape[0] == cov.shape[1],\
                "`cov` must be square matrix"
            assert cov.shape[0] == self.dim and cov.shape[1] == self.dim,\
                "`cov` must have dimension `[dim, dim]`"
            assert _is_symmetric_matrix, "`cov` must be symmetric"
            diag, rot = torch.eig(cov, eigenvectors=True)
            assert torch.allclose(diag[:, 1], torch.zeros_like(diag[:, 1])),\
                "`cov` possesses complex valued eigenvalues"
            diag = diag[:, 0]
            assert torch.all(diag > 0), "`cov` must be positive definite"
            self.register_buffer("_log_diag", diag.log().unsqueeze(0))
            self.register_buffer("_rot", rot)
//  [python]
def __init__(self, dim, mean=None, cov=None):
        super().__init__(dim)
        self._has_mean = mean is not None
        if self._has_mean:
            assert len(mean.shape) == 1,\
                "`mean` must be a vector"
            assert mean.shape[-1] == self.dim, "`mean` must have dimension `dim`"
            self.register_buffer("_mean", mean.unsqueeze(0))
        self._has_cov = cov is not None
        if self._has_cov:
            assert len(cov.shape) == 2 and cov.shape[0] == cov.shape[1],\
                "`cov` must be square matrix"
            assert cov.shape[0] == self.dim and cov.shape[1] == self.dim,\
                "`cov` must have dimension `[dim, dim]`"
            assert _is_symmetric_matrix, "`cov` must be symmetric"
            diag, rot = torch.linalg.eig(cov)
            assert torch.allclose(diag.imag, torch.zeros_like(diag.imag)),\
                "`cov` possesses complex valued eigenvalues"
            diag = diag.real
            assert torch.all(diag > 0), "`cov` must be positive definite"
            self.register_buffer("_log_diag", diag.log().unsqueeze(0))
            self.register_buffer("_rot", rot)
//  [python]
def expm_eig(A):
    eigen_values, eigen_vector = th.eig(A, eigenvectors=True)

    eigen_values.exp_()

    return th.mm(th.mm(eigen_vector, th.diag(eigen_values[:, 0])), eigen_vector.t_())
//  [python]
def expm_eig(A):
    eigen_values, eigen_vector = th.linalg.eig(A)

    eigen_values.exp_()

    return th.mm(th.mm(eigen_vector, th.diag(eigen_values[:, 0])), eigen_vector.t_())
//  [python]
def compute_beta(self,xm):

        M = torch.tensor(xm.size(-2), dtype=torch.float)

        xm = xm.detach().requires_grad_(True)

        Kxx = self.K(xm, xm.detach())

        dKxx_dx = autograd.grad(Kxx.sum(), xm)[0]

        # Kxx = Kxx + eta * I
        if self.eta is not None:
            Kxx += self.eta * torch.eye(xm.size(-2)).to(self.device)

        eigen_vals, eigen_vecs = torch.eig(Kxx, eigenvectors=True)

        if self.num_eigs is not None:
            eigen_vals = eigen_vals[:self.num_eigs]
            eigen_vecs = eigen_vecs[:, :self.num_eigs]



        # Compute the Monte Carlo estimate of the gradient of
        # the eigenfunction at x
        dKxx_dx_avg = -dKxx_dx/xm.shape[0] # [M x D]

        beta = - torch.sqrt(M) * eigen_vecs.t() @ dKxx_dx_avg
        beta *= (1. / eigen_vals[:, 0].unsqueeze(-1))

        return beta, eigen_vals, eigen_vecs
//  [python]
def compute_beta(self,xm):

        M = torch.tensor(xm.size(-2), dtype=torch.float)

        xm = xm.detach().requires_grad_(True)

        Kxx = self.K(xm, xm.detach())

        dKxx_dx = autograd.grad(Kxx.sum(), xm)[0]

        # Kxx = Kxx + eta * I
        if self.eta is not None:
            Kxx += self.eta * torch.eye(xm.size(-2)).to(self.device)

        eigen_vals, eigen_vecs = torch.linalg.eig(Kxx)

        if self.num_eigs is not None:
            eigen_vals = eigen_vals[:self.num_eigs]
            eigen_vecs = eigen_vecs[:, :self.num_eigs]



        # Compute the Monte Carlo estimate of the gradient of
        # the eigenfunction at x
        dKxx_dx_avg = -dKxx_dx/xm.shape[0] # [M x D]

        beta = - torch.sqrt(M) * eigen_vecs.t() @ dKxx_dx_avg
        beta *= (1. / eigen_vals[:, 0].unsqueeze(-1))

        return beta, eigen_vals, eigen_vecs
//  [python]
def _data_eigen_exctractor(self, dataset: List[Graph]) -> pd.DataFrame:
        columns_list = list(map(lambda x: "eigen" + str(x), range(self.n_min)))
        data_bamt = pd.DataFrame(columns=columns_list + ["y"])
        for gr in dataset:
            adj_matrix = to_dense_adj(gr.edge_index)
            eig = torch.eig(adj_matrix.reshape(adj_matrix.shape[1], adj_matrix.shape[2]))[0].T[0].T
            ordered, indices = torch.sort(eig[: gr.num_nodes], descending=True)
            to_append = pd.Series(ordered[: self.n_min].tolist() + gr.y.tolist(), index=data_bamt.columns)
            data_bamt = data_bamt.append(to_append, ignore_index=True)

        return data_bamt
//  [python]
def _data_eigen_exctractor(self, dataset: List[Graph]) -> pd.DataFrame:
        columns_list = list(map(lambda x: "eigen" + str(x), range(self.n_min)))
        data_bamt = pd.DataFrame(columns=columns_list + ["y"])
        for gr in dataset:
            adj_matrix = to_dense_adj(gr.edge_index)
            eig = torch.linalg.eig(adj_matrix.reshape(adj_matrix.shape[1], adj_matrix.shape[2]))[0].T[0].T
            ordered, indices = torch.sort(eig[: gr.num_nodes], descending=True)
            to_append = pd.Series(ordered[: self.n_min].tolist() + gr.y.tolist(), index=data_bamt.columns)
            data_bamt = data_bamt.append(to_append, ignore_index=True)

        return data_bamt
//  [python]
def torch_PCA(X, k, center=True, scale=False):
    X = X.t()
    n,p = X.size()
    ones = torch.ones(n).cuda().view([n,1])
    h = ((1/n) * torch.mm(ones, ones.t())) if center else torch.zeros(n*n).view([n,n])
    H = torch.eye(n).cuda() - h
    X_center =  torch.mm(H.double(), X.double())
    covariance = 1/(n-1) * torch.mm(X_center.t(), X_center).view(p,p)
    scaling = torch.sqrt(1/torch.diag(covariance)).double() if scale else torch.ones(p).cuda().double()
    scaled_covariance = torch.mm(torch.diag(scaling).view(p,p), covariance)
    eigenvalues, eigenvectors = torch.eig(scaled_covariance, True)
    components = (eigenvectors[:, :k])
    #explained_variance = eigenvalues[:k, 0]
    return components
//  [python]
def torch_PCA(X, k, center=True, scale=False):
    X = X.t()
    n,p = X.size()
    ones = torch.ones(n).cuda().view([n,1])
    h = ((1/n) * torch.mm(ones, ones.t())) if center else torch.zeros(n*n).view([n,n])
    H = torch.eye(n).cuda() - h
    X_center =  torch.mm(H.double(), X.double())
    covariance = 1/(n-1) * torch.mm(X_center.t(), X_center).view(p,p)
    scaling = torch.sqrt(1/torch.diag(covariance)).double() if scale else torch.ones(p).cuda().double()
    scaled_covariance = torch.mm(torch.diag(scaling).view(p,p), covariance)
    eigenvalues, eigenvectors = torch.linalg.eig(scaled_covariance)
    components = (eigenvectors[:, :k])
    #explained_variance = eigenvalues[:k, 0]
    return components
//  [python]
def _convolve(dataset: List[Graph], weights: List[float], left_vertices: List[int]) -> List[Graph]:
        new_data = []
        for graph in dataset:
            adj = to_dense_adj(graph.edge_index)
            eigs = torch.eig(adj.reshape(adj.shape[1], adj.shape[2]), True)
            eigenvectors = eigs[1]
            eig = eigs[0].T[0].T
            ordered, indices = torch.sort(eig[: graph.num_nodes], descending=True)
            lef = indices[left_vertices]
            zeroed = torch.tensor(list(set(range(len(eig))) - set(lef.tolist())))
            if len(zeroed) > 0:
                eig[zeroed] = 0

            for e, d in enumerate(lef):
                eig[d] = eig[d] * weights[e]

            eigenvalues = torch.diag(eig)
            convolved = torch.matmul(torch.matmul(eigenvectors, eigenvalues), eigenvectors.T)

            graph.edge_index, graph.edge_weight = dense_to_sparse(convolved)
            graph.edge_weight = graph.edge_weight
            graph.edge_index = graph.edge_index.type(torch.LongTensor)
            new_data.append(graph)
        return new_data
//  [python]
def _convolve(dataset: List[Graph], weights: List[float], left_vertices: List[int]) -> List[Graph]:
        new_data = []
        for graph in dataset:
            adj = to_dense_adj(graph.edge_index)
            eigs = torch.linalg.eig(adj.reshape(adj.shape[1], adj.shape[2]))
            eigenvectors = eigs.eigenvectors
            eig = eigs.eigenvalues.T[0].T
            ordered, indices = torch.sort(eig[: graph.num_nodes], descending=True)
            lef = indices[left_vertices]
            zeroed = torch.tensor(list(set(range(len(eig))) - set(lef.tolist())))
            if len(zeroed) > 0:
                eig[zeroed] = 0

            for e, d in enumerate(lef):
                eig[d] = eig[d] * weights[e]

            eigenvalues = torch.diag(eig)
            convolved = torch.matmul(torch.matmul(eigenvectors, eigenvalues), eigenvectors.T)

            graph.edge_index, graph.edge_weight = dense_to_sparse(convolved)
            graph.edge_weight = graph.edge_weight
            graph.edge_index = graph.edge_index.type(torch.LongTensor)
            new_data.append(graph)
        return new_data
//  [python]
def reset_parameters(self):
        weight_dict = self.state_dict()
        for key, value in weight_dict.items():
            if key == 'weight_ih_l0':
                nn.init.uniform_(value, -1, 1)
                value *= self.w_ih_scale[1:]
            elif re.fullmatch('weight_ih_l[^0]*', key):
                nn.init.uniform_(value, -1, 1)
            elif re.fullmatch('bias_ih_l[0-9]*', key):
                nn.init.uniform_(value, -1, 1)
                value *= self.w_ih_scale[0]
            elif re.fullmatch('weight_hh_l[0-9]*', key):
                w_hh = torch.Tensor(self.hidden_size * self.hidden_size)
                w_hh.uniform_(-1, 1)
                if self.density < 1:
                    zero_weights = torch.randperm(int(self.hidden_size * self.hidden_size))
                    zero_weights = zero_weights[:int(self.hidden_size * self.hidden_size * (1 - self.density))]
                    w_hh[zero_weights] = 0
                w_hh = w_hh.view(self.hidden_size, self.hidden_size)
                abs_eigs = (torch.eig(w_hh)[0] ** 2).sum(1).sqrt()
                weight_dict[key] = w_hh * (self.spectral_radius / torch.max(abs_eigs))
        self.load_state_dict(weight_dict)
//  [python]
def reset_parameters(self):
        weight_dict = self.state_dict()
        for key, value in weight_dict.items():
            if key == 'weight_ih_l0':
                nn.init.uniform_(value, -1, 1)
                value *= self.w_ih_scale[1:]
            elif re.fullmatch('weight_ih_l[^0]*', key):
                nn.init.uniform_(value, -1, 1)
            elif re.fullmatch('bias_ih_l[0-9]*', key):
                nn.init.uniform_(value, -1, 1)
                value *= self.w_ih_scale[0]
            elif re.fullmatch('weight_hh_l[0-9]*', key):
                w_hh = torch.Tensor(self.hidden_size * self.hidden_size)
                w_hh.uniform_(-1, 1)
                if self.density < 1:
                    zero_weights = torch.randperm(int(self.hidden_size * self.hidden_size))
                    zero_weights = zero_weights[:int(self.hidden_size * self.hidden_size * (1 - self.density))]
                    w_hh[zero_weights] = 0
                w_hh = w_hh.view(self.hidden_size, self.hidden_size)
                abs_eigs = (torch.linalg.eigvals(w_hh)**2).abs().sqrt()
                weight_dict[key] = w_hh * (self.spectral_radius / torch.max(abs_eigs))
        self.load_state_dict(weight_dict)
//  [python]
def continuous_lqr(A, B, Q, R, device="cpu"):
    """Solve the continuous time LQR controller for a continuous time system.

    A and B are system matrices, describing the systems dynamics:
     dx/dt = A x + B u

    The controller minimizes the infinite horizon quadratic cost function:
     cost = integral (x.T*Q*x + u.T*R*u) dt

    where Q is a positive semidefinite matrix, and R is positive definite matrix.

    Returns K, X, eigVals:
    Returns gain the optimal gain K, the solution matrix X, and the closed loop system eigenvalues.
    The optimal input is then computed as:
     input: u = -K*x
    """
    # Ref Bertsekas, p.151

    # First, try to solve the continuous Riccati equation
    # NOTE: PyTorch currently not supported by scipy, hence the transfer
    # Need to rework the solver in PyTorch to obtain a speedup
    X = torch.Tensor(solve_continuous_are(
        A.cpu().numpy(), B.cpu().numpy(), Q.cpu().numpy(), R.cpu().numpy())).to(device)

    # Compute the LQR gain
    K = mm(torch.inverse(R), (mm(B.T, X)))

    eigenvalues = torch.eig(A - B * K)

    return K, X, eigenvalues
//  [python]
def continuous_lqr(A, B, Q, R, device="cpu"):
    """Solve the continuous time LQR controller for a continuous time system.

    A and B are system matrices, describing the systems dynamics:
     dx/dt = A x + B u

    The controller minimizes the infinite horizon quadratic cost function:
     cost = integral (x.T*Q*x + u.T*R*u) dt

    where Q is a positive semidefinite matrix, and R is positive definite matrix.

    Returns K, X, eigVals:
    Returns gain the optimal gain K, the solution matrix X, and the closed loop system eigenvalues.
    The optimal input is then computed as:
     input: u = -K*x
    """
    # Ref Bertsekas, p.151

    # First, try to solve the continuous Riccati equation
    # NOTE: PyTorch currently not supported by scipy, hence the transfer
    # Need to rework the solver in PyTorch to obtain a speedup
    X = torch.Tensor(solve_continuous_are(
        A.cpu().numpy(), B.cpu().numpy(), Q.cpu().numpy(), R.cpu().numpy())).to(device)

    # Compute the LQR gain
    K = mm(torch.inverse(R), (mm(B.T, X)))

    eigenvalues = torch.linalg.eig(A - B * K)

    return K, X, eigenvalues
//  [python]
def pca(style_mat):
    # style_mat.shape == (B, d)
    # TODO: naive impl of PCA, should use SVD.
    cov = (style_mat - style_mat.mean(0)).transpose(0, 1) @ (style_mat - style_mat.mean(0))
    eigval, eigvec = torch.eig(cov, eigenvectors=True)
    print(eigval[:, 0])
    return eigvec[:, eigval[:, 0].argsort(descending=True)]
//  [python]
def pca(style_mat):
    # style_mat.shape == (B, d)
    # TODO: naive impl of PCA, should use SVD.
    cov = (style_mat - style_mat.mean(0)).transpose(0, 1) @ (style_mat - style_mat.mean(0))
    eigval, eigvec = torch.linalg.eig(cov)
    print(eigval[:, 0])
    return eigvec[:, eigval[:, 0].argsort(descending=True)]
//  [python]
def NTK(pred):
    paranum = sum(p.numel() for p in model_h.parameters())
    A = torch.zeros((len(pred), paranum)).cuda()
    for index,pred_e in enumerate(pred):
        grad_e = grad(pred_e, model_h.parameters(),retain_graph=True,  create_graph=True) 
        grad_one = torch.tensor([]).cuda() 
        for i in grad_e:
            grad_one = torch.cat((grad_one, i.flatten())) 
        A[index] = grad_one 
    K = torch.mm(A, A.T) 
    L,_ = torch.eig(K.data.cpu()) 
    eig = torch.norm(L, dim=1)
    return eig
//  [python]
def NTK(pred):
    paranum = sum(p.numel() for p in model_h.parameters())
    A = torch.zeros((len(pred), paranum)).cuda()
    for index,pred_e in enumerate(pred):
        grad_e = grad(pred_e, model_h.parameters(),retain_graph=True,  create_graph=True) 
        grad_one = torch.tensor([]).cuda() 
        for i in grad_e:
            grad_one = torch.cat((grad_one, i.flatten())) 
        A[index] = grad_one 
    K = torch.mm(A, A.T) 
    L,_ = torch.linalg.eig(K.data.cpu()) 
    eig = torch.norm(L, dim=1)
    return eig
//  [python]
def discrete_lqr(A, B, Q, R, device="cpu"):
    """Solve the discrete time LQR controller for a discrete time system.

    A and B are system matrices, describing the systems dynamics:
     x[k+1] = A x[k] + B u[k]

    The controller minimizes the infinite horizon quadratic cost function:
     cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]

    where Q is a positive semidefinite matrix, and R is positive definite matrix.

    Returns K, X, eigVals:
    Returns gain the optimal gain K, the solution matrix X, and the closed loop system eigenvalues.
    The optimal input is then computed as:
     input: u = -K*x
    """
    #ref Bertsekas, p.151

    # First, try to solve the discrete Riccati equation
    # NOTE: PyTorch currently not supported by scipy, hence the transfer
    # Need to rework the solver in PyTorch to obtain a speedup
    X = torch.Tensor(solve_discrete_are(
        A.cpu().numpy(), B.cpu().numpy(), Q.cpu().numpy(), R.cpu().numpy())).to(device)

    # Compute the LQR gain
    K = mm(torch.inverse(mm(mm(B.T, X),B)+R), (mm(mm(B.T, X), A)))

    eigenvalues = torch.eig(A - B * K)

    return K, X, eigenvalues
//  [python]
def discrete_lqr(A, B, Q, R, device="cpu"):
    """Solve the discrete time LQR controller for a discrete time system.

    A and B are system matrices, describing the systems dynamics:
     x[k+1] = A x[k] + B u[k]

    The controller minimizes the infinite horizon quadratic cost function:
     cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]

    where Q is a positive semidefinite matrix, and R is positive definite matrix.

    Returns K, X, eigVals:
    Returns gain the optimal gain K, the solution matrix X, and the closed loop system eigenvalues.
    The optimal input is then computed as:
     input: u = -K*x
    """
    #ref Bertsekas, p.151

    # First, try to solve the discrete Riccati equation
    # NOTE: PyTorch currently not supported by scipy, hence the transfer
    # Need to rework the solver in PyTorch to obtain a speedup
    X = torch.Tensor(solve_discrete_are(
        A.cpu().numpy(), B.cpu().numpy(), Q.cpu().numpy(), R.cpu().numpy())).to(device)

    # Compute the LQR gain
    K = mm(torch.inverse(mm(mm(B.T, X),B)+R), (mm(mm(B.T, X), A)))

    eigenvalues = torch.linalg.eig(A - B * K)

    return K, X, eigenvalues
//  [python]
def Dist_Comp(pg1, pg2, dotp):
    A = dotp
    A = A - torch.min(A) #############

    #############################################################################################################3
    b = A.data.clone().contiguous()
    b.requires_grad = False

    # making the diagonal zero
    z = -1 * torch.diag(b)
    kk = Variable(torch.diag(z), requires_grad=False)
    A = kk + A

    B_size = A.size(0)

    F_rank = Variable(torch.zeros(B_size, B_size).cuda())

    for i in range(0, B_size):
        v = torch.ones(B_size, 1)
        v[i] = 0
        ind = (v != 0).nonzero()

        ind = ind[:, 0]
        ind = ind.cuda()
        M = A[ind][:, ind]

        
        M2 = Variable(M.data.clone(), requires_grad=False)
        EE = torch.eig(M2)
        EE2 = EE[0]

        alpha = torch.max(EE2[:, 0])
        alpha = alpha + 0.0005
        # alpha = 0.9

        ident_M = torch.eye(B_size, B_size)
        ident_M[i, i] = 0

        ident_M = ident_M * float(alpha)
        ident_M = Variable(ident_M, requires_grad=False)
        A_2 = A - ident_M.cuda()

        A_2 = A_2 + alpha

        # A_2 = torch.mul(A_2, p_g_score).contiguous()
        # calling Replicator dynamics
        P_g_rank = Replicator(A_2)

        # P_g_rank = torch.exp(P_g_rank)

        # P_g_rank = (P_g_rank - torch.min(P_g_rank))/ (torch.max(P_g_rank) - torch.min(P_g_rank))
        P_g_rank = P_g_rank.view(1, -1)

        F_rank[i, :] = P_g_rank

    # instances_num = 8
    # num_id = B_size / instances_num
    # temp = []

    # rpKnn =
    # pg1=torch.abs(pg1)
    # F_rank = torch.mm(F_rank, p_g_score)
    # pg1 = torch.abs(pg1)

    F_rank_T = (0.3 - F_rank).contiguous()

    F_rank = (0.1 * ((F_rank).cuda()) * (0.9 * (pg1).cuda())).contiguous()
    F_rank2 = (0.1 * (((F_rank_T)).cuda()) * (0.9 * (pg2).cuda())).contiguous()

    # K_nn2 = Variable(K_nn2)
    # Fr = (torch.matmul(F_rank, K_nn2).cuda()).contiguous()  # each column corresponds to the id, and the row corresponds to the gallery image

    return F_rank, F_rank2
//  [python]
def Dist_Comp(pg1, pg2, dotp):
    A = dotp
    A = A - torch.min(A) #############

    #############################################################################################################3
    b = A.data.clone().contiguous()
    b.requires_grad = False

    # making the diagonal zero
    z = -1 * torch.diag(b)
    kk = Variable(torch.diag(z), requires_grad=False)
    A = kk + A

    B_size = A.size(0)

    F_rank = Variable(torch.zeros(B_size, B_size).cuda())

    for i in range(0, B_size):
        v = torch.ones(B_size, 1)
        v[i] = 0
        ind = (v != 0).nonzero()

        ind = ind[:, 0]
        ind = ind.cuda()
        M = A[ind][:, ind]

        
        M2 = Variable(M.data.clone(), requires_grad=False)
        EE = torch.linalg.eig(M2)
        EE2 = EE[0]

        alpha = torch.max(EE2[:, 0])
        alpha = alpha + 0.0005
        # alpha = 0.9

        ident_M = torch.eye(B_size, B_size)
        ident_M[i, i] = 0

        ident_M = ident_M * float(alpha)
        ident_M = Variable(ident_M, requires_grad=False)
        A_2 = A - ident_M.cuda()

        A_2 = A_2 + alpha

        # A_2 = torch.mul(A_2, p_g_score).contiguous()
        # calling Replicator dynamics
        P_g_rank = Replicator(A_2)

        # P_g_rank = torch.exp(P_g_rank)

        # P_g_rank = (P_g_rank - torch.min(P_g_rank))/ (torch.max(P_g_rank) - torch.min(P_g_rank))
        P_g_rank = P_g_rank.view(1, -1)

        F_rank[i, :] = P_g_rank

    # instances_num = 8
    # num_id = B_size / instances_num
    # temp = []

    # rpKnn =
    # pg1=torch.abs(pg1)
    # F_rank = torch.mm(F_rank, p_g_score)
    # pg1 = torch.abs(pg1)

    F_rank_T = (0.3 - F_rank).contiguous()

    F_rank = (0.1 * ((F_rank).cuda()) * (0.9 * (pg1).cuda())).contiguous()
    F_rank2 = (0.1 * (((F_rank_T)).cuda()) * (0.9 * (pg2).cuda())).contiguous()

    # K_nn2 = Variable(K_nn2)
    # Fr = (torch.matmul(F_rank, K_nn2).cuda()).contiguous()  # each column corresponds to the id, and the row corresponds to the gallery image

    return F_rank, F_rank2
//  [python]
def plot_se23_helper(T_est, P_est, v, color, i1, i2):
    """
    Draw ellipse based on the 3 more important directions of the covariance
    """
    D, V = torch.eig(P_est, eigenvectors=True)
    Y, I = torch.sort(D[:, 0], descending=True)
    a = sigma*D[I[0], 0].sqrt() * V[:, I[0]]
    b = sigma*D[I[1], 0].sqrt() * V[:, I[1]]
    c = sigma*D[I[2], 0].sqrt() * V[:, I[2]]
    for n in range(3):
        if n == 0:
            xi = a*v.sin() + b*v.cos()
        elif n == 1:
            xi = b*v.sin() + c*v.cos()
        elif n == 2:
            xi = a*v.sin() + c*v.cos()
        Ttemp = T_est[-1].expand(m_max, 5, 5).bmm(SE3_2.exp(xi.cuda()).cpu())
        clines = Ttemp[:, :3, 4]
        plt.plot(clines[:, i1], clines[:, i2], color=color)
//  [python]
def plot_se23_helper(T_est, P_est, v, color, i1, i2):
    """
    Draw ellipse based on the 3 more important directions of the covariance
    """
    D, V = torch.linalg.eig(P_est)
    Y, I = torch.sort(D[:, 0], descending=True)
    a = sigma*D[I[0], 0].sqrt() * V[:, I[0]]
    b = sigma*D[I[1], 0].sqrt() * V[:, I[1]]
    c = sigma*D[I[2], 0].sqrt() * V[:, I[2]]
    for n in range(3):
        if n == 0:
            xi = a*v.sin() + b*v.cos()
        elif n == 1:
            xi = b*v.sin() + c*v.cos()
        elif n == 2:
            xi = a*v.sin() + c*v.cos()
        Ttemp = T_est[-1].expand(m_max, 5, 5).bmm(SE3_2.exp(xi.cuda()).cpu())
        clines = Ttemp[:, :3, 4]
        plt.plot(clines[:, i1], clines[:, i2], color=color)
//  [python]
def lda(X, y, n_classes, lamb):
    # flatten X
    X = X.view(X.shape[0], -1)
    N, D = X.shape

    # count unique labels in y
    labels, counts = torch.unique(y, return_counts=True)
    assert len(labels) == n_classes  # require X,y cover all classes

    # compute mean-centered observations and covariance matrix
    X_bar = X - torch.mean(X, 0)
    Xc_mean = torch.zeros((n_classes, D), dtype=X.dtype, device=X.device, requires_grad=False)
    St = X_bar.t().matmul(X_bar) / (N - 1)  # total scatter matrix
    Sw = torch.zeros((D, D), dtype=X.dtype, device=X.device, requires_grad=True)  # within-class scatter matrix
    for c, Nc in zip(labels, counts):
        Xc = X[y == c]
        Xc_mean[int(c), :] = torch.mean(Xc, 0)
        Xc_bar = Xc - Xc_mean[int(c), :]
        Sw = Sw + Xc_bar.t().matmul(Xc_bar) / (Nc - 1)
    Sw /= n_classes
    Sb = St - Sw  # between scatter matrix

    # cope for numerical instability
    Sw += torch.eye(D, dtype=X.dtype, device=X.device, requires_grad=False) * lamb

    # compute eigen decomposition
    temp = Sw.pinverse().matmul(Sb)
    # evals, evecs = torch.symeig(temp, eigenvectors=True) # only works for symmetric matrix
    evals, evecs = torch.eig(temp, eigenvectors=True) # shipped from nightly-built version (1.8.0.dev20201015)
    print(evals.shape, evecs.shape)

    # remove complex eigen values and sort
    noncomplex_idx = evals[:, 1] == 0
    evals = evals[:, 0][noncomplex_idx] # take real part of eigen values
    evecs = evecs[:, noncomplex_idx]
    evals, inc_idx = torch.sort(evals) # sort by eigen values, in ascending order
    evecs = evecs[:, inc_idx]
    print(evals.shape, evecs.shape)

    # flag to indicate if to skip backpropagation
    hasComplexEVal = evecs.shape[1] < evecs.shape[0]

    return hasComplexEVal, Xc_mean, evals, evecs
//  [python]
def lda(X, y, n_classes, lamb):
    # flatten X
    X = X.view(X.shape[0], -1)
    N, D = X.shape

    # count unique labels in y
    labels, counts = torch.unique(y, return_counts=True)
    assert len(labels) == n_classes  # require X,y cover all classes

    # compute mean-centered observations and covariance matrix
    X_bar = X - torch.mean(X, 0)
    Xc_mean = torch.zeros((n_classes, D), dtype=X.dtype, device=X.device, requires_grad=False)
    St = X_bar.t().matmul(X_bar) / (N - 1)  # total scatter matrix
    Sw = torch.zeros((D, D), dtype=X.dtype, device=X.device, requires_grad=True)  # within-class scatter matrix
    for c, Nc in zip(labels, counts):
        Xc = X[y == c]
        Xc_mean[int(c), :] = torch.mean(Xc, 0)
        Xc_bar = Xc - Xc_mean[int(c), :]
        Sw = Sw + Xc_bar.t().matmul(Xc_bar) / (Nc - 1)
    Sw /= n_classes
    Sb = St - Sw  # between scatter matrix

    # cope for numerical instability
    Sw += torch.eye(D, dtype=X.dtype, device=X.device, requires_grad=False) * lamb

    # compute eigen decomposition
    temp = Sw.pinverse().matmul(Sb)
    evals, evecs = torch.linalg.eig(temp)
    print(evals.shape, evecs.shape)

    # remove complex eigen values and sort
    noncomplex_idx = evals[:, 1] == 0
    evals = evals[:, 0][noncomplex_idx] # take real part of eigen values
    evecs = evecs[:, noncomplex_idx]
    evals, inc_idx = torch.sort(evals) # sort by eigen values, in ascending order
    evecs = evecs[:, inc_idx]
    print(evals.shape, evecs.shape)

    # flag to indicate if to skip backpropagation
    hasComplexEVal = evecs.shape[1] < evecs.shape[0]

    return hasComplexEVal, Xc_mean, evals, evecs
//  [python]
def compute_w_corr(self, M):
        if self.balance_type == "direct" and self.dyn_convert == 1:
            return M + self.dyn_eps * torch.eye(M.size(0)).to(M.device)

        if self.corr_eigen_decomp:
            if not self.predictor_signaling_2:
                log.info(f"compute_w_corr: Use eigen_decomp of size {M.size()}")
            D, Q = torch.eig(M, eigenvectors=True)
            # Only use the real part. 
            D = D[:,0]
        else:
            # Just use diagonal element. 
            if not self.predictor_signaling_2:
                log.info("compute_w_corr: No eigen_decomp, just use diagonal elements!")
            D = M.diag()
            Q = torch.eye(M.size(0)).to(D.device)
             
        # if eigen_values >= 1, scale everything down. 
        balance_type = self.balance_type
        reg = self.dyn_reg

        if balance_type == "shrink2mean":
            mean_eig = D.mean()
            eigen_values = (D - mean_eig) / 2 + mean_eig
        elif balance_type == "clamp":
            eigen_values = D.clamp(min=0, max=1-reg)
        elif balance_type == "boost_scale":
            max_eig = D.max()
            eigen_values = D.clamp(0) / max_eig
            # Going through a concave function (dyn_convert > 1, e.g., 2 or sqrt function) to boost small eigenvalues (while still keep very small one to be 0)
            if self.dyn_eps_inside:
                # Note that here dyn_eps is allowed to be negative.
                eigen_values = (eigen_values + self.dyn_eps).clamp(1e-4).pow(1/self.dyn_convert)
            else:
                # Note that here dyn_eps is allowed to be negative.
                eigen_values = eigen_values.pow(1/self.dyn_convert) + self.dyn_eps
                eigen_values = eigen_values.clamp(1e-4)
            if not self.predictor_signaling_2:
                sorted_values, _ = eigen_values.sort(descending=True)
                log.info(f"Compute eigenvalues with boost_scale: Top-5: {sorted_values[:5]}, Bottom-5: {sorted_values[-5:]}")
                
            # add regularization to shrink to smaller than 1
            if reg is not None:
                eigen_values *= 1 - reg

        elif balance_type == "scale":
            max_eig = D.max()
            if max_eig > 1 - reg:
                eigen_values = D / (max_eig + reg)
        else:
            raise RuntimeError(f"Unkonwn balance_type: {balance_type}")

        return Q @ eigen_values.diag() @ Q.t()
//  [python]
def compute_w_corr(self, M):
        if self.balance_type == "direct" and self.dyn_convert == 1:
            return M + self.dyn_eps * torch.eye(M.size(0)).to(M.device)

        if self.corr_eigen_decomp:
            if not self.predictor_signaling_2:
                log.info(f"compute_w_corr: Use eigen_decomp of size {M.size()}")
            D, Q = torch.linalg.eig(M)
            # Only use the real part. 
            D = D.real
        else:
            # Just use diagonal element. 
            if not self.predictor_signaling_2:
                log.info("compute_w_corr: No eigen_decomp, just use diagonal elements!")
            D = M.diag()
            Q = torch.eye(M.size(0)).to(D.device)
             
        # if eigen_values >= 1, scale everything down. 
        balance_type = self.balance_type
        reg = self.dyn_reg

        if balance_type == "shrink2mean":
            mean_eig = D.mean()
            eigen_values = (D - mean_eig) / 2 + mean_eig
        elif balance_type == "clamp":
            eigen_values = D.clamp(min=0, max=1-reg)
        elif balance_type == "boost_scale":
            max_eig = D.max()
            eigen_values = D.clamp(0) / max_eig
            # Going through a concave function (dyn_convert > 1, e.g., 2 or sqrt function) to boost small eigenvalues (while still keep very small one to be 0)
            if self.dyn_eps_inside:
                # Note that here dyn_eps is allowed to be negative.
                eigen_values = (eigen_values + self.dyn_eps).clamp(1e-4).pow(1/self.dyn_convert)
            else:
                # Note that here dyn_eps is allowed to be negative.
                eigen_values = eigen_values.pow(1/self.dyn_convert) + self.dyn_eps
                eigen_values = eigen_values.clamp(1e-4)
            if not self.predictor_signaling_2:
                sorted_values, _ = eigen_values.sort(descending=True)
                log.info(f"Compute eigenvalues with boost_scale: Top-5: {sorted_values[:5]}, Bottom-5: {sorted_values[-5:]}")
                
            # add regularization to shrink to smaller than 1
            if reg is not None:
                eigen_values *= 1 - reg

        elif balance_type == "scale":
            max_eig = D.max()
            if max_eig > 1 - reg:
                eigen_values = D / (max_eig + reg)
        else:
            raise RuntimeError(f"Unkonwn balance_type: {balance_type}")

        return Q @ eigen_values.diag() @ Q.t()
//  [python]
def compute_w_minimal_space(self, M, M2, w): 
        try:
            # QDQ^t = M
            D, Q = torch.eig(M, eigenvectors=True)
        except RuntimeError:
            import pdb
            pdb.set_trace()
        # The equation is \dot W = -MW - WM + 2 * M2 (when W keeps symmetric)
        # The solution is e^{-Mt} M2 e^{-Mt}
        # If M can be decomposed: M = QDQ^T, then the solution is Qe^{-Dt}Q^T M2 Q e^{-Dt} Q^T
        d = w.size(1)
        eigen_values = D[:,0].clamp(0)
        M2_convert = Q.t() @ M2 @ Q 
        already_psd = False 
        if self.dyn_time is None:
            # integrate things out and you get a matrix which is 2 / (d_i + d_j), where d_i is the eigenvalues. 
            if self.dyn_diagonalize: 
                M2_diag = M2_convert.diag() / (eigen_values + self.dyn_reg / 2)
                if self.dyn_psd is not None:
                    M2_diag.clamp_(0)
                    already_psd = True
                M2_convert = M2_diag.diag()
            else:
                M2_convert /= (eigen_values.view(d, 1) + eigen_values.view(1, d) + self.dyn_reg) / 2
        else:
            eD = (eigen_values * (-self.dyn_time)).exp().diag()
            M2_convert = eD @ M2_convert @ eD

        w = Q @ M2_convert @ Q.t()

        # Project the weight to be PSD?
        # If we choose to make it PSD (or diagonalized), it will be symmetric by default.
        if self.dyn_psd is None:
            return w
        
        if already_psd:
            w_half = Q @ M2_diag.sqrt().diag()
        else:
            D2, Q2 = torch.eig(w, eigenvectors=True)
            eigen_values2 = D2[:,0].clamp(0)

            # w = w_half @ w_half.t() 
            w_half = Q2 @ eigen_values2.sqrt().diag() 

        if self.dyn_psd > 0:
            # then we do a few iterations to make w more precise.
            alpha = 0.01
            err_magnitudes = torch.zeros(self.dyn_psd)
            w_half_hist = []
            for kk in range(self.dyn_psd):
                w = w_half @ w_half.t()
                err = - (M @ w + w @ M) + M2 * 2
                grad = err @ w_half
                err_magnitudes[kk] = grad.norm()
                w_half_hist.append(w_half.clone())
                # print(f"[{kk}] err: {err.norm().item()}, grad: {grad.norm().item()}")
                w_half += alpha * grad

            err_magnitudes[torch.isnan(err_magnitudes)] = 1e38
            best_kk = err_magnitudes.argmin().item()
            w_half = w_half_hist[best_kk]
        w = w_half @ w_half.t()
        return w
//  [python]
def compute_w_minimal_space(self, M, M2, w): 
        try:
            # QDQ^t = M
            D, Q = torch.linalg.eig(M)
        except RuntimeError:
            import pdb
            pdb.set_trace()
        # The equation is \dot W = -MW - WM + 2 * M2 (when W keeps symmetric)
        # The solution is e^{-Mt} M2 e^{-Mt}
        # If M can be decomposed: M = QDQ^T, then the solution is Qe^{-Dt}Q^T M2 Q e^{-Dt} Q^T
        d = w.size(1)
        eigen_values = D[:,0].clamp(0)
        M2_convert = Q.t() @ M2 @ Q 
        already_psd = False 
        if self.dyn_time is None:
            # integrate things out and you get a matrix which is 2 / (d_i + d_j), where d_i is the eigenvalues. 
            if self.dyn_diagonalize: 
                M2_diag = M2_convert.diag() / (eigen_values + self.dyn_reg / 2)
                if self.dyn_psd is not None:
                    M2_diag.clamp_(0)
                    already_psd = True
                M2_convert = M2_diag.diag()
            else:
                M2_convert /= (eigen_values.view(d, 1) + eigen_values.view(1, d) + self.dyn_reg) / 2
        else:
            eD = (eigen_values * (-self.dyn_time)).exp().diag()
            M2_convert = eD @ M2_convert @ eD

        w = Q @ M2_convert @ Q.t()

        # Project the weight to be PSD?
        # If we choose to make it PSD (or diagonalized), it will be symmetric by default.
        if self.dyn_psd is None:
            return w
        
        if already_psd:
            w_half = Q @ M2_diag.sqrt().diag()
        else:
            D2, Q2 = torch.linalg.eig(w)
            eigen_values2 = D2[:,0].clamp(0)

            # w = w_half @ w_half.t() 
            w_half = Q2 @ eigen_values2.sqrt().diag() 

        if self.dyn_psd > 0:
            # then we do a few iterations to make w more precise.
            alpha = 0.01
            err_magnitudes = torch.zeros(self.dyn_psd)
            w_half_hist = []
            for kk in range(self.dyn_psd):
                w = w_half @ w_half.t()
                err = - (M @ w + w @ M) + M2 * 2
                grad = err @ w_half
                err_magnitudes[kk] = grad.norm()
                w_half_hist.append(w_half.clone())
                # print(f"[{kk}] err: {err.norm().item()}, grad: {grad.norm().item()}")
                w_half += alpha * grad

            err_magnitudes[torch.isnan(err_magnitudes)] = 1e38
            best_kk = err_magnitudes.argmin().item()
            w_half = w_half_hist[best_kk]
        w = w_half @ w_half.t()
        return w
//  [python]
def plot_so3_helper(T_est, P_est, v, color, i1, i2):
    """
    Draw ellipse based on the 3 more important directions of the covariance
    """
    D, V = torch.eig(P_est, eigenvectors=True)
    Y, I = torch.sort(D[:, 0], descending=True)
    a = sigma*D[I[0], 0].sqrt() * V[:, I[0]]
    b = sigma*D[I[1], 0].sqrt() * V[:, I[1]]
    c = sigma*D[I[2], 0].sqrt() * V[:, I[2]]
    for n in range(3):
        if n == 0:
            p = a*v.sin() + b*v.cos()
        elif n == 1:
            p = b*v.sin() + c*v.cos()
        elif n == 2:
            p = a*v.sin() + c*v.cos()
        clines = T_est[-1, :3, 4].expand(m_max, 3) + p
        plt.plot(clines[:, i1], clines[:, i2], color=color)