jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional["LazyTensor"]]:
        """Method that allows implementing special-cased symeig computation. Should not be called directly"""
        from lazy.non_lazy_tensor import NonLazyTensor

        dtype = self.dtype  # perform decomposition in double precision for numerical stability
        # TODO: Use fp64 registry once #1213 is addressed
        evals, evecs = torch.linalg.eigh(self.evaluate().to(dtype=torch.double), eigenvectors=eigenvectors)
        # chop any negative eigenvalues. TODO: warn if evals are significantly negative
        evals = evals.clamp_min(0.0).to(dtype=dtype)
        if eigenvectors:
            evecs = NonLazyTensor(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs
//  [python]
def _compute_kfe(self, group, state):
        """Computes the covariances."""
        mod = group['mod']
        x = self.state[group['mod']]['x']
        gy = self.state[group['mod']]['gy']
        # Computation of xxt
        if group['layer_type'] == 'Conv2d':
            if not self.sua:
                x = F.conv2d(x, group['gathering_filter'],
                             stride=mod.stride, padding=mod.padding,
                             groups=mod.in_channels)
            x = x.data.permute(1, 0, 2, 3).contiguous().view(x.shape[1], -1)
        else:
            x = x.data.t()
        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones], dim=0)
        xxt = torch.mm(x, x.t()) / float(x.shape[1])
        Ex, state['kfe_x'] = torch.symeig(xxt, eigenvectors=True)
        # Computation of ggt
        if group['layer_type'] == 'Conv2d':
            gy = gy.data.permute(1, 0, 2, 3)
            state['num_locations'] = gy.shape[2] * gy.shape[3]
            gy = gy.contiguous().view(gy.shape[0], -1)
        else:
            gy = gy.data.t()
            state['num_locations'] = 1
        ggt = torch.mm(gy, gy.t()) / float(gy.shape[1])
        Eg, state['kfe_gy'] = torch.symeig(ggt, eigenvectors=True)
        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * state['num_locations']
        if group['layer_type'] == 'Conv2d' and self.sua:
            ws = group['params'][0].grad.data.size()
            state['m2'] = state['m2'].view(Eg.size(0), Ex.size(0), 1, 1).expand(-1, -1, ws[2], ws[3])
//  [python]
def _compute_kfe(self, group, state):
        """Computes the covariances."""
        mod = group['mod']
        x = self.state[group['mod']]['x']
        gy = self.state[group['mod']]['gy']
        # Computation of xxt
        if group['layer_type'] == 'Conv2d':
            if not self.sua:
                x = F.conv2d(x, group['gathering_filter'],
                             stride=mod.stride, padding=mod.padding,
                             groups=mod.in_channels)
            x = x.data.permute(1, 0, 2, 3).contiguous().view(x.shape[1], -1)
        else:
            x = x.data.t()
        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones], dim=0)
        xxt = torch.mm(x, x.t()) / float(x.shape[1])
        Ex, state['kfe_x'] = torch.linalg.eigh(xxt)
        # Computation of ggt
        if group['layer_type'] == 'Conv2d':
            gy = gy.data.permute(1, 0, 2, 3)
            state['num_locations'] = gy.shape[2] * gy.shape[3]
            gy = gy.contiguous().view(gy.shape[0], -1)
        else:
            gy = gy.data.t()
            state['num_locations'] = 1
        ggt = torch.mm(gy, gy.t()) / float(gy.shape[1])
        Eg, state['kfe_gy'] = torch.linalg.eigh(ggt)
        state['m2'] = Eg.unsqueeze(1) * Ex.unsqueeze(0) * state['num_locations']
        if group['layer_type'] == 'Conv2d' and self.sua:
            ws = group['params'][0].grad.data.size()
            state['m2'] = state['m2'].view(Eg.size(0), Ex.size(0), 1, 1).expand(-1, -1, ws[2], ws[3])
//  [python]
def compute_hessian(model, dataset):
    loader = torch.utils.data.DataLoader(dataset, batch_size=500)

    n = sum(p.numel() for p in model.parameters())
    h = torch.zeros(n, n, device=device)

    for i, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)

        output = model(data)
        loss = F.nll_loss(output, target, reduction='sum') / len(dataset)

        print('\rCompute the hessian: [{}]'.format("*" * i + " " * (len(loader) - i)), end='')
        hessian(loss, model.parameters(), out=h)
    print('\rCompute the hessian: [{}]'.format("*" * len(loader)))

    print('Diagonalize the hessian...    ')
    eigenvalues, _ = torch.symeig(h)
    print("The first eigenvalues are {}".format(eigenvalues[:20]))
    print("The last eigenvalues are {}".format(eigenvalues[-20:]))
//  [python]
def compute_hessian(model, dataset):
    loader = torch.utils.data.DataLoader(dataset, batch_size=500)

    n = sum(p.numel() for p in model.parameters())
    h = torch.zeros(n, n, device=device)

    for i, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)

        output = model(data)
        loss = F.nll_loss(output, target, reduction='sum') / len(dataset)

        print('\rCompute the hessian: [{}]'.format("*" * i + " " * (len(loader) - i)), end='')
        hessian(loss, model.parameters(), out=h)
    print('\rCompute the hessian: [{}]'.format("*" * len(loader)))

    print('Diagonalize the hessian...    ')
    eigenvalues, _ = torch.linalg.eigh(h)
    print("The first eigenvalues are {}".format(eigenvalues[:20]))
    print("The last eigenvalues are {}".format(eigenvalues[-20:]))
//  [python]
def gdpp_compute_diversity(phi):
    phi = F.normalize(phi, p=2, dim=1)
    SB = torch.mm(phi, phi.t())
    eig_vals, eig_vecs = torch.symeig(SB, eigenvectors=True)

    return eig_vals, eig_vecs
//  [python]
def gdpp_compute_diversity(phi):
    phi = F.normalize(phi, p=2, dim=1)
    SB = torch.mm(phi, phi.t())
    eig_vals, eig_vecs = torch.linalg.eigh(SB)

    return eig_vals, eig_vecs
//  [python]
def DSFA(self, X, Y):

        #m,n = troech.shape(X)
        X_hat = X - torch.mean(X, dim=0)
        Y_hat = Y - torch.mean(Y, dim=0)
        differ = X_hat - Y_hat
        A = torch.matmul(differ.T, differ)
        A = A / self.num

        Sigma_XX = torch.matmul(X_hat.T, X_hat)
        Sigma_XX = Sigma_XX / self.num + self.reg * torch.eye(self.output_num)
        Sigma_YY = torch.matmul(Y_hat.T, Y_hat)
        Sigma_YY = Sigma_YY / self.num + self.reg * torch.eye(self.output_num)
        B = (Sigma_XX + Sigma_YY) / 2

        # For numerical stability.
        D_B, V_B = torch.symeig(B, eigenvectors=True, upper=False) # 只有下三角矩阵参与运算
        #print(D_B.shape, V_B.shape)
        # 取回满足条件索引
        idx = (D_B > 1e-12).nonzero(as_tuple=False)
        idx = idx[:, 0]
        # 实现tf.gather
        D_B = torch.gather(D_B, 0, idx)
        idx_vb = idx.unsqueeze(1)
        idx_vb = idx_vb.expand(1, idx_vb.shape[0], V_B.shape[1])
        idx_vb = idx_vb.squeeze()
        V_B = torch.gather(V_B, 1, idx_vb)
        B_inv = torch.matmul(torch.matmul(V_B, torch.diag(torch.reciprocal(D_B))), V_B.T)

        Sigma = torch.matmul(B_inv, A)
        loss = torch.trace(torch.matmul(Sigma, Sigma))

        return loss
//  [python]
def DSFA(self, X, Y):

        X_hat = X - torch.mean(X, dim=0)
        Y_hat = Y - torch.mean(Y, dim=0)
        differ = X_hat - Y_hat
        A = torch.matmul(differ.T, differ)
        A = A / self.num

        Sigma_XX = torch.matmul(X_hat.T, X_hat)
        Sigma_XX = Sigma_XX / self.num + self.reg * torch.eye(self.output_num)
        Sigma_YY = torch.matmul(Y_hat.T, Y_hat)
        Sigma_YY = Sigma_YY / self.num + self.reg * torch.eye(self.output_num)
        B = (Sigma_XX + Sigma_YY) / 2

        # For numerical stability.
        D_B, V_B = torch.linalg.eigh(B)
        #print(D_B.shape, V_B.shape)
        # 取回满足条件索引
        idx = (D_B > 1e-12).nonzero(as_tuple=False)
        idx = idx[:, 0]
        # 实现tf.gather
        D_B = torch.gather(D_B, 0, idx)
        idx_vb = idx.unsqueeze(1)
        idx_vb = idx_vb.expand(1, idx_vb.shape[0], V_B.shape[1])
        idx_vb = idx_vb.squeeze()
        V_B = torch.gather(V_B, 1, idx_vb)
        B_inv = torch.matmul(torch.matmul(V_B, torch.diag(torch.reciprocal(D_B))), V_B.T)

        Sigma = torch.matmul(B_inv, A)
        loss = torch.trace(torch.matmul(Sigma, Sigma))

        return loss
//  [python]
def test(self, test_loader):
        self.model.eval()
        
        with torch.no_grad():

            landmark_embeddings = []

            for landmark_batch in self.landmark_loader:
                landmark_batch_embeddings = self.model(landmark_batch)
                landmark_embeddings = landmark_embeddings + landmark_batch_embeddings

            self.gamma = self.compute_gamma(landmark_embeddings) # no backpropagation for gamma
            
            if self.nystrom == "LLSVM":
                
                K_Z = self.compute_mmd_gram_matrix(landmark_embeddings).to(self.device)
                eigenvalues, U_Z = torch.symeig(K_Z, eigenvectors=True)

                #removed smallest 2/3 eigenvalues due to numerical instability
                no_of_eigens = len(eigenvalues)
                eigenvalues = eigenvalues[-no_of_eigens//3:]
                
                # if eigenvalues still negative, adjust - values small enough so that it does not affect
                m = min(eigenvalues)
                if m < 0:
                    print("neg")
                    eigenvalues = eigenvalues - 2*m
                elif m == 0:
                    print("zero")
                    eigenvalues = eigenvalues + 1e-9
                
                U_Z = U_Z[:,-no_of_eigens//3:]
                T = torch.matmul(U_Z,torch.diag(eigenvalues**-0.5))
            
            dists_list = []
            for batch in test_loader:

                R_embeddings = self.model(batch)
                K_RZ = self.compute_mmd_gram_matrix(R_embeddings, landmark_embeddings).to(self.device)
                
                if self.nystrom == "LLSVM":
                    F = torch.matmul(K_RZ, T)
                elif self.nystrom == "RSVM":
                    F = K_RZ
                
                batch_dists = torch.sum((F - self.center)**2, dim=1).cpu()
                dists_list.append(batch_dists)
            
            labels = torch.cat([batch.y for batch in test_loader])
            dists = torch.cat(dists_list)
            
            ap = average_precision_score(labels, dists)
            roc_auc = roc_auc_score(labels, dists)

            return ap, roc_auc, dists, labels
//  [python]
def test(self, test_loader):
        self.model.eval()
        
        with torch.no_grad():

            landmark_embeddings = []

            for landmark_batch in self.landmark_loader:
                landmark_batch_embeddings = self.model(landmark_batch)
                landmark_embeddings = landmark_embeddings + landmark_batch_embeddings

            self.gamma = self.compute_gamma(landmark_embeddings) # no backpropagation for gamma
            
            if self.nystrom == "LLSVM":
                
                K_Z = self.compute_mmd_gram_matrix(landmark_embeddings).to(self.device)
                eigenvalues, U_Z = torch.linalg.eigh(K_Z)

                #removed smallest 2/3 eigenvalues due to numerical instability
                no_of_eigens = len(eigenvalues)
                eigenvalues = eigenvalues[-no_of_eigens//3:]
                
                # if eigenvalues still negative, adjust - values small enough so that it does not affect
                m = min(eigenvalues)
                if m < 0:
                    print("neg")
                    eigenvalues = eigenvalues - 2*m
                elif m == 0:
                    print("zero")
                    eigenvalues = eigenvalues + 1e-9
                
                U_Z = U_Z[:,-no_of_eigens//3:]
                T = torch.matmul(U_Z,torch.diag(eigenvalues**-0.5))
            
            dists_list = []
            for batch in test_loader:

                R_embeddings = self.model(batch)
                K_RZ = self.compute_mmd_gram_matrix(R_embeddings, landmark_embeddings).to(self.device)
                
                if self.nystrom == "LLSVM":
                    F = torch.matmul(K_RZ, T)
                elif self.nystrom == "RSVM":
                    F = K_RZ
                
                batch_dists = torch.sum((F - self.center)**2, dim=1).cpu()
                dists_list.append(batch_dists)
            
            labels = torch.cat([batch.y for batch in test_loader])
            dists = torch.cat(dists_list)
            
            ap = average_precision_score(labels, dists)
            roc_auc = roc_auc_score(labels, dists)

            return ap, roc_auc, dists, labels
//  [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 train(self, train_loader):
        self.model.train()
        
        if self.center == None:  # first iteration
            F_list = []

        loss_accum = 0
        svdd_loss_accum = 0
        reg_loss_accum = 0
        total_iters = 0

        if self.debug_mode:
            torch.autograd.set_detect_anomaly(True)

        for batch in train_loader:

            landmark_embeddings = []
            for landmark_batch in self.landmark_loader:
                landmark_batch_embeddings = self.model(landmark_batch)
                landmark_embeddings = landmark_embeddings + landmark_batch_embeddings

            self.gamma = self.compute_gamma(landmark_embeddings).detach() # no backpropagation for gamma
            
            train_embeddings = self.model(batch)
            K_trainZ = self.compute_mmd_gram_matrix(train_embeddings, landmark_embeddings).to(self.device)

            if self.nystrom == "LLSVM":
                
                K_Z = self.compute_mmd_gram_matrix(landmark_embeddings).to(self.device)

                K_temp = K_Z.detach()
                eps_matrix = torch.randn_like(K_Z)*torch.median(torch.abs(K_temp))*1e-4
                eigenvalues, U_Z = torch.symeig(K_Z+eps_matrix,eigenvectors=True)

                #removed smallest 2/3 eigenvalues due to numerical instability
                no_of_eigens = len(eigenvalues)
                eigenvalues = eigenvalues[-no_of_eigens//3:]
                
                # if eigenvalues still negative, adjust - values small enough so that it does not affect
                m = min(eigenvalues).detach()
                if m < 0:
                    eigenvalues = eigenvalues - 2*m
                elif m == 0:
                    eigenvalues = eigenvalues + 1e-9
                
                U_Z = U_Z[:,-no_of_eigens//3:]
                T = torch.matmul(U_Z,torch.diag(eigenvalues**-0.5))
                F_train = torch.matmul(K_trainZ, T)

            elif self.nystrom == "RSVM":
                
                F_train = K_trainZ
                
            

            # if first iteration, compute center, and don't do any backprop
            if self.center == None:
                F_list.append(F_train)
            
            else:
                train_scores = torch.sum((F_train - self.center)**2, dim=1).cpu()
                svdd_loss = torch.mean(train_scores)
                
                #backpropagate
                self.optimizer.zero_grad()
                svdd_loss.backward()
                self.optimizer.step()
                
                svdd_loss_accum += svdd_loss.detach().cpu().numpy()
                total_iters += 1

        if self.center == None:
            full_F_list = torch.cat(F_list)
            self.center = torch.mean(full_F_list, dim=0).detach() # no backpropagation for center
            #print("center computed")

            average_svdd_loss = -1
        
        else:
            average_svdd_loss = svdd_loss_accum/total_iters

        return average_svdd_loss
//  [python]
def train(self, train_loader):
        self.model.train()
        
        if self.center == None:  # first iteration
            F_list = []

        loss_accum = 0
        svdd_loss_accum = 0
        reg_loss_accum = 0
        total_iters = 0

        if self.debug_mode:
            torch.autograd.set_detect_anomaly(True)

        for batch in train_loader:

            landmark_embeddings = []
            for landmark_batch in self.landmark_loader:
                landmark_batch_embeddings = self.model(landmark_batch)
                landmark_embeddings = landmark_embeddings + landmark_batch_embeddings

            self.gamma = self.compute_gamma(landmark_embeddings).detach() # no backpropagation for gamma
            
            train_embeddings = self.model(batch)
            K_trainZ = self.compute_mmd_gram_matrix(train_embeddings, landmark_embeddings).to(self.device)

            if self.nystrom == "LLSVM":
                
                K_Z = self.compute_mmd_gram_matrix(landmark_embeddings).to(self.device)

                K_temp = K_Z.detach()
                eps_matrix = torch.randn_like(K_Z)*torch.median(torch.abs(K_temp))*1e-4
                eigenvalues, U_Z = torch.linalg.eigh(K_Z+eps_matrix)

                #removed smallest 2/3 eigenvalues due to numerical instability
                no_of_eigens = len(eigenvalues)
                eigenvalues = eigenvalues[-no_of_eigens//3:]
                
                # if eigenvalues still negative, adjust - values small enough so that it does not affect
                m = min(eigenvalues).detach()
                if m < 0:
                    eigenvalues = eigenvalues - 2*m
                elif m == 0:
                    eigenvalues = eigenvalues + 1e-9
                
                U_Z = U_Z[:,-no_of_eigens//3:]
                T = torch.matmul(U_Z,torch.diag(eigenvalues**-0.5))
                F_train = torch.matmul(K_trainZ, T)

            elif self.nystrom == "RSVM":
                
                F_train = K_trainZ
                
            

            # if first iteration, compute center, and don't do any backprop
            if self.center == None:
                F_list.append(F_train)
            
            else:
                train_scores = torch.sum((F_train - self.center)**2, dim=1).cpu()
                svdd_loss = torch.mean(train_scores)
                
                #backpropagate
                self.optimizer.zero_grad()
                svdd_loss.backward()
                self.optimizer.step()
                
                svdd_loss_accum += svdd_loss.detach().cpu().numpy()
                total_iters += 1

        if self.center == None:
            full_F_list = torch.cat(F_list)
            self.center = torch.mean(full_F_list, dim=0).detach() # no backpropagation for center
            #print("center computed")

            average_svdd_loss = -1
        
        else:
            average_svdd_loss = svdd_loss_accum/total_iters

        return average_svdd_loss
//  [python]
def get_entropy(self) -> to.Tensor:
        """
        Get the exploration distribution's entropy.
        The entropy of a normal distribution is independent of the mean.

        :return: entropy value
        """
        # det_cov = to.cholesky(self.cov).diag().prod()  # cholesky returns lower triangular matrix
        det_cov = to.prod(to.symeig(self.cov.data, eigenvectors=False)[0])
        return self.cov.shape[0] / 2.0 * (1.0 + to.log(to.tensor(2 * math.pi))) + to.log(det_cov) / 2.0
//  [python]
def get_entropy(self) -> to.Tensor:
        """
        Get the exploration distribution's entropy.
        The entropy of a normal distribution is independent of the mean.

        :return: entropy value
        """
        # det_cov = to.cholesky(self.cov).diag().prod()  # cholesky returns lower triangular matrix
        det_cov = to.prod(to.linalg.eigh(self.cov.data, eigenvectors=False)[0])
        return self.cov.shape[0] / 2.0 * (1.0 + to.log(to.tensor(2 * math.pi))) + to.log(det_cov) / 2.0
//  [python]
def step(self, closure=None):
        # 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.0, 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, closure=None):
        # 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.0, 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 create_latent_frame(pnts,n_nodes):
    pnts = pnts.transpose(1,2)
    center = pnts.mean(2,True)

    pnts_centered = pnts - center
    
    R = torch.bmm(pnts_centered,pnts_centered.transpose(1,2))
    lambdas,V_ = torch.symeig(R.detach().cpu(),True)
    F =  V_.to(R)
    ops = torch.tensor([[1,1,1],
                        [1,1,-1],
                        [1,-1,1],
                        [1,-1,-1],
                        [-1,1,1],
                        [-1,1,-1],
                        [-1,-1,1],
                        [-1,-1,-1]]).unsqueeze(1).to(F)
    F_ops = ops.unsqueeze(0) * F.unsqueeze(1)
    framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts - center).transpose(1,2))
    # framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts).transpose(1,2))
    framed_input = framed_input.transpose(0,1)
    framed_input = framed_input.view(framed_input.shape[0],framed_input.shape[1],n_nodes, -1, 3)
    framed_input = torch.reshape(framed_input,(-1,framed_input.shape[-2]*3))
    return framed_input, F_ops.detach(), center.detach()
//  [python]
def create_latent_frame(pnts,n_nodes):
    pnts = pnts.transpose(1,2)
    center = pnts.mean(2,True)

    pnts_centered = pnts - center
    
    R = torch.bmm(pnts_centered,pnts_centered.transpose(1,2))
    lambdas,V_ = torch.linalg.eigh(R.detach().cpu())
    F =  V_.to(R)
    ops = torch.tensor([[1,1,1],
                        [1,1,-1],
                        [1,-1,1],
                        [1,-1,-1],
                        [-1,1,1],
                        [-1,1,-1],
                        [-1,-1,1],
                        [-1,-1,-1]]).unsqueeze(1).to(F)
    F_ops = ops.unsqueeze(0) * F.unsqueeze(1)
    framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts - center).transpose(1,2))
    # framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts).transpose(1,2))
    framed_input = framed_input.transpose(0,1)
    framed_input = framed_input.view(framed_input.shape[0],framed_input.shape[1],n_nodes, -1, 3)
    framed_input = torch.reshape(framed_input,(-1,framed_input.shape[-2]*3))
    return framed_input, F_ops.detach(), center.detach()
//  [python]
def loss(self, *args):
        # https: // www.uta.edu / math / _docs / preprint / 2014 / rep2014_04.pdf
        # H is n_views * n_samples * k
        all_views = [view - view.mean(dim=0) for view in args]

        eigen_views = [view @ torch.inverse(view.T @ view) @ view.T for view in all_views]

        Q = torch.stack(eigen_views, dim=0).sum(dim=0)

        [eigvals, eigvecs] = torch.symeig(Q, eigenvectors=True)

        idx = torch.argsort(eigvals, descending=True)
        eigvecs = eigvecs[:, idx]

        corr = (eigvals[idx][:self.outdim_size] - 1).sum()

        return -corr
//  [python]
def loss(self, *args):
        # https: // www.uta.edu / math / _docs / preprint / 2014 / rep2014_04.pdf
        # H is n_views * n_samples * k
        all_views = [view - view.mean(dim=0) for view in args]

        eigen_views = [view @ torch.inverse(view.T @ view) @ view.T for view in all_views]

        Q = torch.stack(eigen_views, dim=0).sum(dim=0)

        [eigvals, eigvecs] = torch.linalg.eigh(Q)

        idx = torch.argsort(eigvals, descending=True)
        eigvecs = eigvecs[:, idx]

        corr = (eigvals[idx][:self.outdim_size] - 1).sum()

        return -corr
//  [python]
def create_frame(nodes, n_nodes):
    pnts = nodes[:,:3]
    v = nodes[:,3:]
    pnts = pnts.view(-1, n_nodes, 3).transpose(1,2)
    v = v.view(-1, n_nodes, 3).transpose(1,2)
    center = pnts.mean(2,True)
    pnts_centered = pnts - center
    # add noise
    R = torch.bmm(pnts_centered,pnts_centered.transpose(1,2))
    lambdas,V_ = torch.symeig(R.detach().cpu(),True)
    F =  V_.to(R)
    ops = torch.tensor([[1,1,1],
                        [1,1,-1],
                        [1,-1,1],
                        [1,-1,-1],
                        [-1,1,1],
                        [-1,1,-1],
                        [-1,-1,1],
                        [-1,-1,-1]]).unsqueeze(1).to(F)
    F_ops = ops.unsqueeze(0) * F.unsqueeze(1)
    framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts - center).transpose(1,2))
    framed_v = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(v).transpose(1,2))
    framed_input = framed_input.transpose(0,1)
    framed_input = torch.reshape(framed_input,(-1,3))
    framed_v = framed_v.transpose(0,1)
    framed_v = torch.reshape(framed_v,(-1,3))
    out = torch.cat([framed_input,framed_v],dim=1)
    return out, F_ops.detach(), center.detach()
//  [python]
def create_frame(nodes, n_nodes):
    pnts = nodes[:,:3]
    v = nodes[:,3:]
    pnts = pnts.view(-1, n_nodes, 3).transpose(1,2)
    v = v.view(-1, n_nodes, 3).transpose(1,2)
    center = pnts.mean(2,True)
    pnts_centered = pnts - center
    # add noise
    R = torch.bmm(pnts_centered,pnts_centered.transpose(1,2))
    lambdas,V_ = torch.linalg.eigh(R.detach().cpu())
    F =  V_.to(R)
    ops = torch.tensor([[1,1,1],
                        [1,1,-1],
                        [1,-1,1],
                        [1,-1,-1],
                        [-1,1,1],
                        [-1,1,-1],
                        [-1,-1,1],
                        [-1,-1,-1]]).unsqueeze(1).to(F)
    F_ops = ops.unsqueeze(0) * F.unsqueeze(1)
    framed_input = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(pnts - center).transpose(1,2))
    framed_v = torch.einsum('boij,bpj->bopi',F_ops.transpose(2,3),(v).transpose(1,2))
    framed_input = framed_input.transpose(0,1)
    framed_input = torch.reshape(framed_input,(-1,3))
    framed_v = framed_v.transpose(0,1)
    framed_v = torch.reshape(framed_v,(-1,3))
    out = torch.cat([framed_input,framed_v],dim=1)
    return out, F_ops.detach(), center.detach()
//  [python]
def identity_PCA(x, dim=3):
    """Reduce dimension of x based on PCA on the union of the n-simplex.
    This is a way of reducing the dimension of x while treating all
    input dimensions with the same importance, independently of the
    input distribution in x.
    """
    assert x.dim() == 2, f"Expected x.dim()=2 but got x.dim()={x.dim()} instead"

    # Create z the union of the N-simplex
    input_dim = x.shape[1]
    z = torch.eye(input_dim)

    # PCA on z
    z_offset = z.mean(axis=0)
    z_centered = z - z_offset
    cov_matrix = z_centered.T.mm(z_centered) / len(z_centered)
    _, eigenvectors = torch.symeig(cov_matrix, eigenvectors=True)

    # Apply the PCA on x
    x_reduced = (x - z_offset).mm(eigenvectors[:, -dim:])

    return x_reduced
//  [python]
def identity_PCA(x, dim=3):
    """Reduce dimension of x based on PCA on the union of the n-simplex.
    This is a way of reducing the dimension of x while treating all
    input dimensions with the same importance, independently of the
    input distribution in x.
    """
    assert x.dim() == 2, f"Expected x.dim()=2 but got x.dim()={x.dim()} instead"

    # Create z the union of the N-simplex
    input_dim = x.shape[1]
    z = torch.eye(input_dim)

    # PCA on z
    z_offset = z.mean(axis=0)
    z_centered = z - z_offset
    cov_matrix = z_centered.T.mm(z_centered) / len(z_centered)
    _, eigenvectors = torch.linalg.eigh(cov_matrix)

    # Apply the PCA on x
    x_reduced = (x - z_offset).mm(eigenvectors[:, -dim:])

    return x_reduced
//  [python]
def adapt(self, mean: to.Tensor = None, cov: Union[to.Tensor, float] = None):
        """
        Adapt the mean and the variance of the noise on the action or parameters.
        Use `None` to leave one of the parameters at their current value.

        :param mean: exploration strategy's new mean
        :param cov: exploration strategy's new standard deviation
        """
        if not (isinstance(mean, to.Tensor) or mean is None):
            raise pyrado.TypeErr(given=mean, expected_type=to.Tensor)
        if not (isinstance(cov, to.Tensor) or cov is None):
            raise pyrado.TypeErr(given=cov, expected_type=to.Tensor)

        if mean is not None:
            assert self.mean is not None, "Can not change fixed zero mean!"
            if not mean.shape == self.mean.shape:
                raise pyrado.ShapeErr(given=mean, expected_match=self.mean)
            self.mean.data = mean

        if cov is not None:
            if not cov.shape == self.cov.shape:
                raise pyrado.ShapeErr(given=cov, expected_match=self.cov)
            self.cov.data = cov

            # Check the eigenvalues of the new matrix and make if positive definite if necessary
            eigs = to.symeig(self.cov.data, eigenvectors=False)[0]
            if not min(eigs > 0.0):
                repair_diag = min(eigs - 1e-3) * to.eye(self.cov.shape[0])
                self.cov.data -= repair_diag  # shifting the eigenvalues
                warn("The covariance matrix has not been positive definite.", UserWarning)

            # Check if the diagonal elements of the new covariance matrix are neither too small nor too big
            for i in range(cov.shape[0]):
                self.cov.data[i, i] = to.max(self.cov.data[i, i], self.std_min**2)
//  [python]
def adapt(self, mean: to.Tensor = None, cov: Union[to.Tensor, float] = None):
        """
        Adapt the mean and the variance of the noise on the action or parameters.
        Use `None` to leave one of the parameters at their current value.

        :param mean: exploration strategy's new mean
        :param cov: exploration strategy's new standard deviation
        """
        if not (isinstance(mean, to.Tensor) or mean is None):
            raise pyrado.TypeErr(given=mean, expected_type=to.Tensor)
        if not (isinstance(cov, to.Tensor) or cov is None):
            raise pyrado.TypeErr(given=cov, expected_type=to.Tensor)

        if mean is not None:
            assert self.mean is not None, "Can not change fixed zero mean!"
            if not mean.shape == self.mean.shape:
                raise pyrado.ShapeErr(given=mean, expected_match=self.mean)
            self.mean.data = mean

        if cov is not None:
            if not cov.shape == self.cov.shape:
                raise pyrado.ShapeErr(given=cov, expected_match=self.cov)
            self.cov.data = cov

            # Check the eigenvalues of the new matrix and make if positive definite if necessary
            eigs = to.linalg.eigh(self.cov.data, eigenvectors=False)[0]
            if not min(eigs > 0.0):
                repair_diag = min(eigs - 1e-3) * to.eye(self.cov.shape[0])
                self.cov.data -= repair_diag  # shifting the eigenvalues
                warn("The covariance matrix has not been positive definite.", UserWarning)

            # Check if the diagonal elements of the new covariance matrix are neither too small nor too big
            for i in range(cov.shape[0]):
                self.cov.data[i, i] = to.max(self.cov.data[i, i], self.std_min**2)
//  [python]
def teig(x):
    e, v = torch.symeig(x, eigenvectors=True)
    return e, v
    try:
        v, e, _ = torch.svd(x)
    except:
        __import__("pdb").set_trace()
    return e, v
//  [python]
def teig(x):
    e, v = torch.linalg.eigh(x)
    return e, v
    try:
        v, e, _ = torch.svd(x)
    except:
        __import__("pdb").set_trace()
    return e, v
//  [python]
def forward(ctx, X):
        if len(X.shape) == 2:
            assert X.shape[0] == X.shape[1]
            assert (X == X.t()).all(), "X must be symmetric."
            Lam, Q = torch.symeig(X, eigenvectors=True)
            output = Q @ torch.diag(torch.exp(Lam)) @ Q.T
        elif len(X.shape) == 3:
            assert X.shape[1] == X.shape[2]
            assert (X == X.transpose(1, 2)).all(), "X must be symmetric."
            Lam, Q = torch.symeig(X, eigenvectors=True)
            output = torch.bmm(torch.bmm(Q, torch.diag_embed(torch.exp(Lam))), Q.transpose(1, 2))
        else:
            raise ArgumentError("X must be 2 or 3 dimensional.")
        ctx.save_for_backward(X, Lam, Q)
            
        return output
//  [python]
def forward(ctx, X):
        if len(X.shape) == 2:
            assert X.shape[0] == X.shape[1]
            assert (X == X.t()).all(), "X must be symmetric."
            Lam, Q = torch.linalg.eigh(X)
            output = Q @ torch.diag(torch.exp(Lam)) @ Q.T
        elif len(X.shape) == 3:
            assert X.shape[1] == X.shape[2]
            assert (X == X.transpose(1, 2)).all(), "X must be symmetric."
            Lam, Q = torch.linalg.eigh(X)
            output = torch.bmm(torch.bmm(Q, torch.diag_embed(torch.exp(Lam))), Q.transpose(1, 2))
        else:
            raise ArgumentError("X must be 2 or 3 dimensional.")
        ctx.save_for_backward(X, Lam, Q)
            
        return output
//  [python]
def sym_l2_norm(mat: torch.Tensor):
    """Largest eigenvalue assuming that matrix is symmetric."""

    u.check_symmetric(mat)

    if gl.debug_linalg_crashes:
        try:
            evals, _evecs = torch.symeig(mat)
        except RuntimeError as e:
            print(e)
            dump(mat, '/tmp/sym_l2_norm.txt')
            sys.exit()
    else:
        evals, _evecs = torch.symeig(mat)

    return torch.max(evals)
//  [python]
def sym_l2_norm(mat: torch.Tensor):
    """Largest eigenvalue assuming that matrix is symmetric."""

    u.check_symmetric(mat)

    if gl.debug_linalg_crashes:
        try:
            evals, _evecs = torch.linalg.eigh(mat)
        except RuntimeError as e:
            print(e)
            dump(mat, '/tmp/sym_l2_norm.txt')
            sys.exit()
    else:
        evals, _evecs = torch.linalg.eigh(mat)

    return torch.max(evals)
//  [python]
def lyapunov_spectral(A, B, cond=None):
    u.check_symmetric(A)
    u.check_symmetric(B)

    s, U = torch.symeig(A, eigenvectors=True)
    if cond is None:
        cond = get_condition(s.dtype)
    cutoff = cond * max(s)
    s = torch.where(s > cutoff, s, torch.tensor(0.).to(s.device))

    C = U.t() @ B @ U    # TODO(y): throw away eigenvectors corresponding to discarded evals. U=U[:num_eigs]
    s = s.unsqueeze(1) + s.unsqueeze(0)
    si = torch.where(s > 0, 1 / s, s)
    Y = C * si
    X = U @ Y @ U.t()

    # cancel small asymetries introduces by multiplication by small numbers
    X = (X + X.t()) / 2
    return X
//  [python]
def lyapunov_spectral(A, B, cond=None):
    u.check_symmetric(A)
    u.check_symmetric(B)

    s, U = torch.linalg.eigh(A)
    if cond is None:
        cond = get_condition(s.dtype)
    cutoff = cond * max(s)
    s = torch.where(s > cutoff, s, torch.tensor(0.).to(s.device))

    C = U.t() @ B @ U    # TODO(y): throw away eigenvectors corresponding to discarded evals. U=U[:num_eigs]
    s = s.unsqueeze(1) + s.unsqueeze(0)
    si = torch.where(s > 0, 1 / s, s)
    Y = C * si
    X = U @ Y @ U.t()

    # cancel small asymetries introduces by multiplication by small numbers
    X = (X + X.t()) / 2
    return X
//  [python]
def lyapunov_svd(A, C, rtol=1e-4, eps=1e-7, use_svd=False):
    """Solve AX+XA=C using SVD"""

    # This method doesn't work for singular matrices, so regularize it
    # TODO: can optimize performance by reusing eigenvalues from regularization computations
    A = regularize_mat(A, eps)
    C = regularize_mat(C, eps)

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = robust_svd(A)
    else:
        S, U = torch.symeig(A, eigenvectors=True)
    S = S.diag() @ torch.ones_like(A)
    X = U @ ((U.t() @ C @ U) / (S + S.t())) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        # TODO(y): currently spams with errors, implement another method based on Newton iteration
        pass
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X
//  [python]
def lyapunov_svd(A, C, rtol=1e-4, eps=1e-7, use_svd=False):
    """Solve AX+XA=C using SVD"""

    # This method doesn't work for singular matrices, so regularize it
    # TODO: can optimize performance by reusing eigenvalues from regularization computations
    A = regularize_mat(A, eps)
    C = regularize_mat(C, eps)

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = robust_svd(A)
    else:
        S, U = torch.linalg.eigh(A)
    S = S.diag() @ torch.ones_like(A)
    X = U @ ((U.t() @ C @ U) / (S + S.t())) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        # TODO(y): currently spams with errors, implement another method based on Newton iteration
        pass
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X
//  [python]
def symeig_pos_evals(mat: torch.Tensor) -> torch.Tensor:
    """Returns positive eigenvalues from symeig in decreasing order (to match order of .svd())"""

    s, u = torch.symeig(mat, eigenvectors=False)
    return torch.flip(filter_evals(s, remove_negative=True), dims=[0])
//  [python]
def symeig_pos_evals(mat: torch.Tensor) -> torch.Tensor:
    """Returns positive eigenvalues from symeig in decreasing order (to match order of .svd())"""

    s, u = torch.linalg.eigh(mat)
    return torch.flip(filter_evals(s, remove_negative=True), dims=[0])
//  [python]
def loss(self, H1, H2):
        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)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + self.r * torch.eye(o1, dtype=torch.double)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + self.r * torch.eye(o2, dtype=torch.double)

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

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

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

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

        # just the top self.outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        # ensures positive definite
        U, V = torch.symeig(trace_TT, eigenvectors=True)
        U_inds = torch.gt(D1, self.eps).nonzero()[:, 0]
        U = U[U_inds]
        corr = torch.sum(torch.sqrt(U))

        return -corr
//  [python]
def loss(self, H1, H2):
        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)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + self.r * torch.eye(o1, dtype=torch.double)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + self.r * torch.eye(o2, dtype=torch.double)

        [D1, V1] = torch.linalg.eigh(SigmaHat11)
        [D2, V2] = torch.linalg.eigh(SigmaHat22)

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

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

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

        # just the top self.outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        # ensures positive definite
        U, V = torch.linalg.eigh(trace_TT)
        U_inds = torch.gt(D1, self.eps).nonzero()[:, 0]
        U = U[U_inds]
        corr = torch.sum(torch.sqrt(U))

        return -corr
//  [python]
def lyapunov_svd2(A, C, rtol=1e-4, eps=1e-7, use_svd=False):
    """Solve AX+XA=C using SVD"""

    # This method doesn't work for singular matrices, so regularize it
    # TODO: can optimize performance by reusing eigenvalues from regularization computations
    # A = regularize_mat(A, eps)
    # C = regularize_mat(C, eps)

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = robust_svd(A)
    else:
        S, U = torch.symeig(A, eigenvectors=True)
    S = S.diag() @ torch.ones_like(A)
    factor = (S + S.t())
    cutoff = max(S) * get_condition(S)
    factor = torch.where(factor > cutoff, 1 / factor, factor)
    X = U @ ((U.t() @ C @ U) * factor) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        # TODO(y): currently spams with errors, implement another method based on Newton iteration
        pass
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X