jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input:
            - M:      [bs, num_corr, num_corr] the compatibility matrix
            - method: select different method for calculating the learding eigenvector.
        Output:
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.symeig(M, eigenvectors=True)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        """
        Calculate the leading eigenvector using power iteration algorithm or torch.symeig
        Input:
            - M:      [bs, num_corr, num_corr] the compatibility matrix
            - method: select different method for calculating the learding eigenvector.
        Output:
            - solution: [bs, num_corr] leading eigenvector
        """
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.linalg.eigh(M)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def train(model, optimizer, subjects_adj,subjects_labels, args):
  
  i = 0
  all_epochs_loss = []
  no_epochs = args.epochs

  for epoch in range(no_epochs):

      epoch_loss = []
      epoch_error = []

      for lr,hr in zip(subjects_adj,subjects_labels):

          model.train()
          optimizer.zero_grad()
          
          lr = torch.from_numpy(lr).type(torch.FloatTensor)
          hr = torch.from_numpy(hr).type(torch.FloatTensor)
          
          model_outputs,net_outs,start_gcn_outs,layer_outs = model(lr)
          model_outputs  = unpad(model_outputs, args.padding)

          padded_hr = pad_HR_adj(hr,args.padding)
          eig_val_hr, U_hr = torch.symeig(padded_hr, eigenvectors=True,upper=True)
          
          loss = args.lmbda * criterion(net_outs, start_gcn_outs) + criterion(model.layer.weights,U_hr) + criterion(model_outputs, hr) 
          
          error = criterion(model_outputs, hr)
          
          loss.backward()
          optimizer.step()

          epoch_loss.append(loss.item())
          epoch_error.append(error.item())
      
      i+=1
      print("Epoch: ",i, "Loss: ", np.mean(epoch_loss), "Error: ", np.mean(epoch_error)*100,"%")
      all_epochs_loss.append(np.mean(epoch_loss))
//  [python]
def train(model, optimizer, subjects_adj,subjects_labels, args):
  
  i = 0
  all_epochs_loss = []
  no_epochs = args.epochs

  for epoch in range(no_epochs):

      epoch_loss = []
      epoch_error = []

      for lr,hr in zip(subjects_adj,subjects_labels):

          model.train()
          optimizer.zero_grad()
          
          lr = torch.from_numpy(lr).type(torch.FloatTensor)
          hr = torch.from_numpy(hr).type(torch.FloatTensor)
          
          model_outputs,net_outs,start_gcn_outs,layer_outs = model(lr)
          model_outputs  = unpad(model_outputs, args.padding)

          padded_hr = pad_HR_adj(hr,args.padding)
          eig_val_hr, U_hr = torch.linalg.eigh(padded_hr)
          
          loss = args.lmbda * criterion(net_outs, start_gcn_outs) + criterion(model.layer.weights,U_hr) + criterion(model_outputs, hr) 
          
          error = criterion(model_outputs, hr)
          
          loss.backward()
          optimizer.step()

          epoch_loss.append(loss.item())
          epoch_error.append(error.item())
      
      i+=1
      print("Epoch: ",i, "Loss: ", np.mean(epoch_loss), "Error: ", np.mean(epoch_error)*100,"%")
      all_epochs_loss.append(np.mean(epoch_loss))
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

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

            la = self.damping + self.weight_decay

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

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

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

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

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

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

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

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

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

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

            la = self.damping + self.weight_decay

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

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

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

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

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

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

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

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

        self.optim.step()
        self.steps += 1
//  [python]
def CCA_loss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

    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()) + r1 * torch.eye(o1, device=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

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

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

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

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

    if use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.symeig(trace_TT, eigenvectors=True)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def CCA_loss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

    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()) + r1 * torch.eye(o1, device=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

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

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

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

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

    if use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.linalg.eigh(trace_TT)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def jacobian_eig(x, eig_idx, model, layer=-1):
    "Get gradient for images"
    x.requires_grad_(True)
    hiddens, _ = model.bothOutputs(x)
    temp2 = hiddens[layer] - torch.mean(hiddens[layer], 0)
    print(temp2.shape)
    cov = temp2.t() @ temp2 / temp2.shape[0]
    cov = (cov + cov.t()) / 2
    eigs, _ = torch.symeig(cov, eigenvectors=True)
    print(eigs.shape)
    gradient = torch.autograd.grad(eigs[-1 + eig_idx], x)[0]
    return gradient
//  [python]
def jacobian_eig(x, eig_idx, model, layer=-1):
    "Get gradient for images"
    x.requires_grad_(True)
    hiddens, _ = model.bothOutputs(x)
    temp2 = hiddens[layer] - torch.mean(hiddens[layer], 0)
    print(temp2.shape)
    cov = temp2.t() @ temp2 / temp2.shape[0]
    cov = (cov + cov.t()) / 2
    eigs, _ = torch.linalg.eigh(cov)
    print(eigs.shape)
    gradient = torch.autograd.grad(eigs[-1 + eig_idx], x)[0]
    return gradient
//  [python]
def mat_root_inv_for_debug(A):
    D, U=torch.symeig(A, eigenvectors=True)
    return U@torch.diag(1/torch.sqrt(D))@U.t()
//  [python]
def mat_root_inv_for_debug(A):
    D, U=torch.linalg.eigh(A)
    return U@torch.diag(1/torch.sqrt(D))@U.t()
//  [python]
def compute_eig_vectors(x, y, model, loss, device):
    hidden, outputs = model.bothOutputs(x.to(device))
    loss = loss(outputs, y.to(device))
    spectraTemp = []
    eigVec = []
    regul = torch.zeros(1, device=device)

    for idx in range(len(hidden)):
        hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        # _, eigTemp, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
        eigTemp, vecTemp = torch.symeig(cov, eigenvectors=True)
        eig_T = eigTemp.float()
        vecTemp = vecTemp.float()
        eigTemp, rT = eigen_val_regulate(0, 0, eig_T, device)  # compute regularizer
        regul += rT
        spectraTemp.append(eigTemp.cpu())  # save spectra
        eigVec.append(vecTemp)

    return eigVec, loss, spectraTemp, regul.cpu().item()
//  [python]
def compute_eig_vectors(x, y, model, loss, device):
    hidden, outputs = model.bothOutputs(x.to(device))
    loss = loss(outputs, y.to(device))
    spectraTemp = []
    eigVec = []
    regul = torch.zeros(1, device=device)

    for idx in range(len(hidden)):
        hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        eigTemp, vecTemp = torch.linalg.eigh(cov)
        eig_T = eigTemp.float()
        vecTemp = vecTemp.float()
        eigTemp, rT = eigen_val_regulate(0, 0, eig_T, device)  # compute regularizer
        regul += rT
        spectraTemp.append(eigTemp.cpu())  # save spectra
        eigVec.append(vecTemp)

    return eigVec, loss, spectraTemp, regul.cpu().item()
//  [python]
def compute_eig_vectors_only(hidden, only_last=False):
    # import pdb; pdb.set_trace()
    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        _, vecTemp = torch.symeig(cov, eigenvectors=True)
        eigVec = vecTemp.float().cpu()
        # torch.pca_lowrank(hTemp)
    else:
        eigVec = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            _, vecTemp = torch.symeig(cov, eigenvectors=True)
            vecTemp = vecTemp.float()
            eigVec.append(vecTemp.cpu())

    return eigVec
//  [python]
def compute_eig_vectors_only(hidden, only_last=False):
    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        _, vecTemp = torch.linalg.eigh(cov)
        eigVec = vecTemp.float().cpu()
        # torch.pca_lowrank(hTemp)
    else:
        eigVec = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            _, vecTemp = torch.linalg.eigh(cov)
            vecTemp = vecTemp.float()
            eigVec.append(vecTemp.cpu())

    return eigVec
//  [python]
def eig_stable(A, tol=1e-10, out=None): 
    #L, V = torch.eig(A, eigenvectors=True)
    #L=L[:,0].view(-1)
    
    L, V = torch.symeig(A, eigenvectors=True)
    keep = L > tol
    #print(out)
    #print("A",A)
    print(keep)
    return(L[keep], V[:,keep])
//  [python]
def eig_stable(A, tol=1e-10, out=None): 
    L, V = torch.linalg.eigh(A)
    keep = L > tol
    return(L[keep], V[:,keep])
//  [python]
def compute_eig_values_only(hidden, only_last=False):

    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        eigTemp, _ = torch.symeig(cov)
        eigVals = eigTemp.float()
    else:
        eigVals = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            eigTemp, _ = torch.symeig(cov)
            eigTemp = eigTemp.float()
            eigVals.append(eigTemp.cpu())
    return eigVals
//  [python]
def compute_eig_values_only(hidden, only_last=False):

    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        eigTemp = torch.linalg.eigh(cov).eigenvalues
        eigVals = eigTemp.float()
    else:
        eigVals = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            eigTemp = torch.linalg.eigh(cov).eigenvalues
            eigTemp = eigTemp.float()
            eigVals.append(eigTemp.cpu())
    return eigVals
//  [python]
def eigens(patches):
    n,c,h,w = patches.shape
    Σ = cov(patches.reshape(n, c*h*w))
    Λ, V = torch.symeig(Σ, eigenvectors=True)
    return Λ.flip(0), V.t().reshape(c*h*w, c, h, w).flip(0)
//  [python]
def eigens(patches):
    n,c,h,w = patches.shape
    Σ = cov(patches.reshape(n, c*h*w))
    Λ, V = torch.linalg.eigh(Σ)
    return Λ.flip(0), V.t().reshape(c*h*w, c, h, w).flip(0)
//  [python]
def update_bays_reg_BDQN(self, log):
        with torch.no_grad():
            log.debug("######## Start updating bayesian linear layer ########")
            # Update action posterior with formulas: \beta | z,y ~ N(mu_q, cov_q)
            z = self.latent_z.double()
            y = self.train_y.squeeze().cuda()
            s = torch.matmul(z.T, z)
            A = s / self.sigma_n + 1 / self.sigma * self.eye
            B = torch.matmul(z.T, y.double()) / self.sigma_n
            # A_eig_val, A_eig_vec = torch.symeig(A, eigenvectors=True)
            inv = torch.inverse(A.double())
            self.mu_w[0] = torch.matmul(inv, B).squeeze()
            temp_cov = self.sigma * inv
            eig_val, eig_vec = torch.symeig(temp_cov, eigenvectors=True)
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[:20]) )
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[-20:]) )
            if torch.any(eig_val < 0):
                self.cov_w[0] = torch.matmul(torch.matmul(eig_vec, torch.diag(torch.abs(eig_val))), torch.t(eig_vec))
            else:
                self.cov_w[0] = temp_cov
//  [python]
def update_bays_reg_BDQN(self, log):
        with torch.no_grad():
            log.debug("######## Start updating bayesian linear layer ########")
            # Update action posterior with formulas: \beta | z,y ~ N(mu_q, cov_q)
            z = self.latent_z.double()
            y = self.train_y.squeeze().cuda()
            s = torch.matmul(z.T, z)
            A = s / self.sigma_n + 1 / self.sigma * self.eye
            B = torch.matmul(z.T, y.double()) / self.sigma_n
            # A_eig_val, A_eig_vec = torch.symeig(A, eigenvectors=True)
            inv = torch.inverse(A.double())
            self.mu_w[0] = torch.matmul(inv, B).squeeze()
            temp_cov = self.sigma * inv
            eig_val, eig_vec = torch.linalg.eigh(temp_cov)
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[:20]) )
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[-20:]) )
            if torch.any(eig_val < 0):
                self.cov_w[0] = torch.matmul(torch.matmul(eig_vec, torch.diag(torch.abs(eig_val))), torch.t(eig_vec))
            else:
                self.cov_w[0] = temp_cov
//  [python]
def __init__(self, random_seed=None):
		self.n_vertices = np.array([PESTCONTROL_N_CHOICE] * PESTCONTROL_N_STAGES)
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = str(random_seed).zfill(4)
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [python]
def __init__(self, random_seed=None):
		self.n_vertices = np.array([PESTCONTROL_N_CHOICE] * PESTCONTROL_N_STAGES)
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = str(random_seed).zfill(4)
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.linalg.eigh(laplacian)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [python]
def __init__(self, A, layer = 5, gamma = 0.9):
        
        super(SGNN, self).__init__()
        
        self.gamma = gamma
        
        self.layer = layer
        
        self.n = A.shape[0]
        A = torch.tensor(A, dtype = dtype, device = device)
        D_rsqrt = torch.tensor(torch.diag(torch.sum(A, dim=0).pow(-0.5)), dtype = dtype, device = device)
        I = torch.tensor(torch.eye(self.n, self.n), dtype = dtype, device = device)
        self.L = I - D_rsqrt.mm(A).mm(D_rsqrt)
        self.e, self.v = torch.symeig(self.L, eigenvectors=True)
        
        self.F_x_param = Parameter(torch.Tensor(self.n))
        self.F_h_param = Parameter(torch.Tensor(self.n))
        
        self.A1 = Parameter(torch.Tensor(self.n))
        self.A2 = Parameter(torch.Tensor(self.n))
        self.A3 = Parameter(torch.Tensor(self.n))
        self.A4 = Parameter(torch.Tensor(self.n))
        self.A5 = Parameter(torch.Tensor(self.n))
        self.A6 = Parameter(torch.Tensor(self.n))
        self.A7 = Parameter(torch.Tensor(self.n))
        self.A8 = Parameter(torch.Tensor(self.n))
        self.A9 = Parameter(torch.Tensor(self.n))
        self.A10 = Parameter(torch.Tensor(self.n))

        self.reset_parameters()
//  [python]
def __init__(self, A, layer = 5, gamma = 0.9):
        
        super(SGNN, self).__init__()
        
        self.gamma = gamma
        
        self.layer = layer
        
        self.n = A.shape[0]
        A = torch.tensor(A, dtype = dtype, device = device)
        D_rsqrt = torch.tensor(torch.diag(torch.sum(A, dim=0).pow(-0.5)), dtype = dtype, device = device)
        I = torch.tensor(torch.eye(self.n, self.n), dtype = dtype, device = device)
        self.L = I - D_rsqrt.mm(A).mm(D_rsqrt)
        self.e, self.v = torch.linalg.eigh(self.L)
        
        self.F_x_param = Parameter(torch.Tensor(self.n))
        self.F_h_param = Parameter(torch.Tensor(self.n))
        
        self.A1 = Parameter(torch.Tensor(self.n))
        self.A2 = Parameter(torch.Tensor(self.n))
        self.A3 = Parameter(torch.Tensor(self.n))
        self.A4 = Parameter(torch.Tensor(self.n))
        self.A5 = Parameter(torch.Tensor(self.n))
        self.A6 = Parameter(torch.Tensor(self.n))
        self.A7 = Parameter(torch.Tensor(self.n))
        self.A8 = Parameter(torch.Tensor(self.n))
        self.A9 = Parameter(torch.Tensor(self.n))
        self.A10 = Parameter(torch.Tensor(self.n))

        self.reset_parameters()
//  [python]
def __init__(self, random_seed_pair=(None, None)):
		self.n_vertices = np.array([CENTROID_N_CHOICE] * CENTROID_N_EDGES)
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
		self.interaction_list = []
		self.covariance_list = []
		self.partition_original_list = []
		self.n_ising_models = 3
		ising_seeds = np.random.RandomState(random_seed_pair[0]).randint(0, 10000, (self.n_ising_models,))
		for i in range(self.n_ising_models):
			interaction = generate_ising_interaction(CENTROID_GRID[0], CENTROID_GRID[1], ising_seeds[i])
			interaction = (interaction[0].numpy(), interaction[1].numpy())
			covariance, partition_original = spin_covariance(interaction, CENTROID_GRID)
			self.interaction_list.append(interaction)
			self.covariance_list.append(covariance)
			self.partition_original_list.append(partition_original)
//  [python]
def __init__(self, random_seed_pair=(None, None)):
		self.n_vertices = np.array([CENTROID_N_CHOICE] * CENTROID_N_EDGES)
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.linalg.eigh(laplacian)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
		self.interaction_list = []
		self.covariance_list = []
		self.partition_original_list = []
		self.n_ising_models = 3
		ising_seeds = np.random.RandomState(random_seed_pair[0]).randint(0, 10000, (self.n_ising_models,))
		for i in range(self.n_ising_models):
			interaction = generate_ising_interaction(CENTROID_GRID[0], CENTROID_GRID[1], ising_seeds[i])
			interaction = (interaction[0].numpy(), interaction[1].numpy())
			covariance, partition_original = spin_covariance(interaction, CENTROID_GRID)
			self.interaction_list.append(interaction)
			self.covariance_list.append(covariance)
			self.partition_original_list.append(partition_original)
//  [python]
def spectral_clustering(sim_matrix: Tensor, cluster_num: int, init: Tensor=None,
                        return_state: bool=False, normalized: bool=False):
    r"""
    Perform spectral clustering based on given similarity matrix.

    This function firstly computes the leading eigenvectors of the given similarity matrix, and then utilizes the
    eigenvectors as features and performs k-means clustering based on these features.

    :param sim_matrix: :math:`(n\times n)` input similarity matrix. :math:`n`: number of instances
    :param cluster_num: number of clusters
    :param init: the initialization technique or initial features for k-means
    :param return_state: whether return state features (can be further used for prediction)
    :param normalized: whether to normalize the similarity matrix by its degree
    :return: the belonging of each instance to clusters, state features (if ``return_state==True``)
    """
    degree = torch.diagflat(torch.sum(sim_matrix, dim=-1))
    if normalized:
        aff_matrix = (degree - sim_matrix) / torch.diag(degree).unsqueeze(1)
    else:
        aff_matrix = degree - sim_matrix
    e, v = torch.symeig(aff_matrix, eigenvectors=True)
    topargs = torch.argsort(torch.abs(e), descending=False)[1:cluster_num]
    v = v[:, topargs]

    if cluster_num == 2:
        choice_cluster = (v > 0).to(torch.int).squeeze(1)
    else:
        choice_cluster, initial_state = kmeans(v, cluster_num, init_x=init if init is not None else 'plus',
                                               distance='euclidean', tol=1e-6)

    choice_cluster = choice_cluster.to(sim_matrix.device)

    if return_state:
        return choice_cluster, initial_state
    else:
        return choice_cluster
//  [python]
def spectral_clustering(sim_matrix: Tensor, cluster_num: int, init: Tensor=None,
                        return_state: bool=False, normalized: bool=False):
    r"""
    Perform spectral clustering based on given similarity matrix.

    This function firstly computes the leading eigenvectors of the given similarity matrix, and then utilizes the
    eigenvectors as features and performs k-means clustering based on these features.

    :param sim_matrix: :math:`(n\times n)` input similarity matrix. :math:`n`: number of instances
    :param cluster_num: number of clusters
    :param init: the initialization technique or initial features for k-means
    :param return_state: whether return state features (can be further used for prediction)
    :param normalized: whether to normalize the similarity matrix by its degree
    :return: the belonging of each instance to clusters, state features (if ``return_state==True``)
    """
    degree = torch.diagflat(torch.sum(sim_matrix, dim=-1))
    if normalized:
        aff_matrix = (degree - sim_matrix) / torch.diag(degree).unsqueeze(1)
    else:
        aff_matrix = degree - sim_matrix
    e, v = torch.linalg.eigh(aff_matrix)
    topargs = torch.argsort(torch.abs(e), descending=False)[1:cluster_num]
    v = v[:, topargs]

    if cluster_num == 2:
        choice_cluster = (v > 0).to(torch.int).squeeze(1)
    else:
        choice_cluster, initial_state = kmeans(v, cluster_num, init_x=init if init is not None else 'plus',
                                               distance='euclidean', tol=1e-6)

    choice_cluster = choice_cluster.to(sim_matrix.device)

    if return_state:
        return choice_cluster, initial_state
    else:
        return choice_cluster
//  [python]
def train(model, subjects_adj, subjects_labels, args):

    bce_loss = nn.BCELoss()
    netD = Discriminator(args)
    print(netD)
    optimizerG = optim.Adam(model.parameters(), lr=args.lr)
    optimizerD = optim.Adam(netD.parameters(), lr=args.lr)

    all_epochs_loss = []
    for epoch in range(args.epochs):
        with torch.autograd.set_detect_anomaly(True):
            epoch_loss = []
            epoch_error = []
            for lr, hr in zip(subjects_adj, subjects_labels):
                optimizerD.zero_grad()
                optimizerG.zero_grad()

                hr = pad_HR_adj(hr, args.padding)
                lr = torch.from_numpy(lr).type(torch.FloatTensor)
                padded_hr = torch.from_numpy(hr).type(torch.FloatTensor)

                eig_val_hr, U_hr = torch.symeig(
                    padded_hr, eigenvectors=True, upper=True)

                model_outputs, net_outs, start_gcn_outs, layer_outs = model(
                    lr, args.lr_dim, args.hr_dim)

                mse_loss = args.lmbda * criterion(net_outs, start_gcn_outs) + criterion(
                    model.layer.weights, U_hr) + criterion(model_outputs, padded_hr)

                error = criterion(model_outputs, padded_hr)
                real_data = model_outputs.detach()
                fake_data = gaussian_noise_layer(padded_hr, args)

                d_real = netD(real_data)
                d_fake = netD(fake_data)

                dc_loss_real = bce_loss(d_real, torch.ones(args.hr_dim, 1))
                dc_loss_fake = bce_loss(d_fake, torch.zeros(args.hr_dim, 1))
                dc_loss = dc_loss_real + dc_loss_fake

                dc_loss.backward()
                optimizerD.step()

                d_fake = netD(gaussian_noise_layer(padded_hr, args))

                gen_loss = bce_loss(d_fake, torch.ones(args.hr_dim, 1))
                generator_loss = gen_loss + mse_loss
                generator_loss.backward()
                optimizerG.step()

                epoch_loss.append(generator_loss.item())
                epoch_error.append(error.item())

            print("Epoch: ", epoch, "Loss: ", np.mean(epoch_loss),
                  "Error: ", np.mean(epoch_error)*100, "%")
            all_epochs_loss.append(np.mean(epoch_loss))
//  [python]
def train(model, subjects_adj, subjects_labels, args):

    bce_loss = nn.BCELoss()
    netD = Discriminator(args)
    print(netD)
    optimizerG = optim.Adam(model.parameters(), lr=args.lr)
    optimizerD = optim.Adam(netD.parameters(), lr=args.lr)

    all_epochs_loss = []
    for epoch in range(args.epochs):
        with torch.autograd.set_detect_anomaly(True):
            epoch_loss = []
            epoch_error = []
            for lr, hr in zip(subjects_adj, subjects_labels):
                optimizerD.zero_grad()
                optimizerG.zero_grad()

                hr = pad_HR_adj(hr, args.padding)
                lr = torch.from_numpy(lr).type(torch.FloatTensor)
                padded_hr = torch.from_numpy(hr).type(torch.FloatTensor)

                eig_val_hr, U_hr = torch.linalg.eigh(
                    padded_hr, eigenvectors=True, upper=True)

                model_outputs, net_outs, start_gcn_outs, layer_outs = model(
                    lr, args.lr_dim, args.hr_dim)

                mse_loss = args.lmbda * criterion(net_outs, start_gcn_outs) + criterion(
                    model.layer.weights, U_hr) + criterion(model_outputs, padded_hr)

                error = criterion(model_outputs, padded_hr)
                real_data = model_outputs.detach()
                fake_data = gaussian_noise_layer(padded_hr, args)

                d_real = netD(real_data)
                d_fake = netD(fake_data)

                dc_loss_real = bce_loss(d_real, torch.ones(args.hr_dim, 1))
                dc_loss_fake = bce_loss(d_fake, torch.zeros(args.hr_dim, 1))
                dc_loss = dc_loss_real + dc_loss_fake

                dc_loss.backward()
                optimizerD.step()

                d_fake = netD(gaussian_noise_layer(padded_hr, args))

                gen_loss = bce_loss(d_fake, torch.ones(args.hr_dim, 1))
                generator_loss = gen_loss + mse_loss
                generator_loss.backward()
                optimizerG.step()

                epoch_loss.append(generator_loss.item())
                epoch_error.append(error.item())

            print("Epoch: ", epoch, "Loss: ", np.mean(epoch_loss),
                  "Error: ", np.mean(epoch_error)*100, "%")
            all_epochs_loss.append(np.mean(epoch_loss))
//  [python]
def evd(
        tensor: torch.Tensor,
        reg_value: float = 0.,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs eigendecomposition of a symmetric hermitian tensor. The eigenvalues
    of the nearest positive semidefinit matrix are returned to ensure
    differentiability.

    Parameters
    ----------
    tensor : :py:class:`torch.Tensor` (nx, nx)
        This tensor is eigen decomposed. The values of the nearest positive
        semidefinit matrix to this tensor are returned.
    reg_value : float, optional
        This regularization value is added to the eigenvalues and represents
        a regularization of the matrix diagonal. The regularization is
        deactivated with the default value of 0.

    Returns
    -------
    evals : :py:class:`torch.Tensor` (nx)
        The eigenvalues of the nearest positive semidefinit matrix to the
        given tensor. The regularization value is already added to these
        eigenvalues.
    evects : :py:class:`torch.Tensor` (nx, nx)
        The estimated eigenvectors based on given tensor.
    evals_inv : :py:class:`torch.Tensor` (nx)
        The inverted eigenvalues of the nearest positive semidefinit matrix
        to the given tensor.
    """
    evals, evects = torch.symeig(tensor, eigenvectors=True, upper=False)
    evals = evals.clamp(min=0)
    evals = evals + reg_value
    evals_inv = 1 / evals
    return evals, evects, evals_inv
//  [python]
def evd(
        tensor: torch.Tensor,
        reg_value: float = 0.,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs eigendecomposition of a symmetric hermitian tensor. The eigenvalues
    of the nearest positive semidefinit matrix are returned to ensure
    differentiability.

    Parameters
    ----------
    tensor : :py:class:`torch.Tensor` (nx, nx)
        This tensor is eigen decomposed. The values of the nearest positive
        semidefinit matrix to this tensor are returned.
    reg_value : float, optional
        This regularization value is added to the eigenvalues and represents
        a regularization of the matrix diagonal. The regularization is
        deactivated with the default value of 0.

    Returns
    -------
    evals : :py:class:`torch.Tensor` (nx)
        The eigenvalues of the nearest positive semidefinit matrix to the
        given tensor. The regularization value is already added to these
        eigenvalues.
    evects : :py:class:`torch.Tensor` (nx, nx)
        The estimated eigenvectors based on given tensor.
    evals_inv : :py:class:`torch.Tensor` (nx)
        The inverted eigenvalues of the nearest positive semidefinit matrix
        to the given tensor.
    """
    evals, evects = torch.linalg.eigh(tensor)
    evals = evals.clamp(min=0)
    evals = evals + reg_value
    evals_inv = 1 / evals
    return evals, evects, evals_inv
//  [python]
def E0_torchAD(model):
    """
        Compute various order of derivatives of the ground state energy using AD
    of the full eigensolver built in Pytorch.
    """
    Es, psis = torch.symeig(model.Hmatrix, eigenvectors=True)
    E0 = Es[0]
    dE0, = torch.autograd.grad(E0, model.g, create_graph=True)
    d2E0, = torch.autograd.grad(dE0, model.g, retain_graph=True)
    return E0.item() / model.N, \
           dE0.item() / model.N, \
           d2E0.item() / model.N
//  [python]
def E0_torchAD(model):
    """
        Compute various order of derivatives of the ground state energy using AD
    of the full eigensolver built in Pytorch.
    """
    Es, psis = torch.linalg.eigh(model.Hmatrix)
    E0 = Es[0]
    dE0, = torch.autograd.grad(E0, model.g, create_graph=True)
    d2E0, = torch.autograd.grad(dE0, model.g, retain_graph=True)
    return E0.item() / model.N, \
           dE0.item() / model.N, \
           d2E0.item() / model.N
//  [python]
def _break_down_cov(self, cov):
        vals, vecs = torch.symeig(cov.to(torch.float64), eigenvectors=True)
        sqrt_vals = vals.sqrt()
        scale = sqrt_vals.unsqueeze(-1) * vecs.t()  # PCA whitening
        inv_scale = vecs / sqrt_vals
        log_sqrt_vals = vals.log().sum().view((1, 1)) / 2
        return scale, inv_scale, log_sqrt_vals
//  [python]
def _break_down_cov(self, cov):
        vals, vecs = torch.linalg.eigh(cov.to(torch.float64))
        sqrt_vals = vals.sqrt()
        scale = sqrt_vals.unsqueeze(-1) * vecs.t()  # PCA whitening
        inv_scale = vecs / sqrt_vals
        log_sqrt_vals = vals.log().sum().view((1, 1)) / 2
        return scale, inv_scale, log_sqrt_vals
//  [python]
def ortho_loss(A, eps=1e-5):
    I = torch.eye(3).cuda()
    epsI = I*eps
    C = torch.mul(A.transpose(2, 1), A) + epsI
    # print(C[0])
    s = torch.symeig(C, eigenvectors=True)[0]
    s1, s2, s3 = s[:, 0], s[:, 1], s[:, 2]

    def elem_sym_polys_of_eigen_values(M):
            M = [[M[:, i, j] for j in range(3)] for i in range(3)]
            sigma1 = M[0][0]+ M[1][1] + M[2][2]

            sigma2 = ((M[0][0] * M[1][1]
                     + M[1][1] * M[2][2]
                     + M[2][2] * M[0][0])
                     -
                      (M[0][1] * M[1][0]
                     + M[1][2] * M[2][1]
                     + M[2][0] * M[0][2]))

            sigma3 = ((M[0][0] * M[1][1] * M[2][2]
                     + M[0][1] * M[1][2] * M[2][0]
                     + M[0][2] * M[1][0] * M[2][1])
                     -
                      (M[0][0] * M[1][2] * M[2][1]
                     + M[0][1] * M[1][0] * M[2][2]
                     + M[0][2] * M[1][1] * M[2][0]))

            return sigma1, sigma2, sigma3

    # sigma1, sigma2, sigma3 represent a+b+c, ab+ac+bc, abc respectively
    # s1, s2, s3 = elem_sym_polys_of_eigen_values(C)
    # print(s1, s2, s3)

    # ortho_loss = s1 + (1 + eps) * (1 + eps) * s2 / s3 - 3 * 2 * (1 + eps)

    ortho_loss = (s1 + (1 + eps) * (1 + eps) / s1) \
               + (s2 + (1 + eps) * (1 + eps) / s2) \
               + (s3 + (1 + eps) * (1 + eps) / s3) - 2 * 3 * (1 + eps)

    return torch.mean(ortho_loss)
//  [python]
def ortho_loss(A, eps=1e-5):
    I = torch.eye(3).cuda()
    epsI = I*eps
    C = torch.mul(A.transpose(2, 1), A) + epsI
    # print(C[0])
    s = torch.linalg.eigh(C, eigenvectors=True)[0]
    s1, s2, s3 = s[:, 0], s[:, 1], s[:, 2]

    def elem_sym_polys_of_eigen_values(M):
            M = [[M[:, i, j] for j in range(3)] for i in range(3)]
            sigma1 = M[0][0]+ M[1][1] + M[2][2]

            sigma2 = ((M[0][0] * M[1][1]
                     + M[1][1] * M[2][2]
                     + M[2][2] * M[0][0])
                     -
                      (M[0][1] * M[1][0]
                     + M[1][2] * M[2][1]
                     + M[2][0] * M[0][2]))

            sigma3 = ((M[0][0] * M[1][1] * M[2][2]
                     + M[0][1] * M[1][2] * M[2][0]
                     + M[0][2] * M[1][0] * M[2][1])
                     -
                      (M[0][0] * M[1][2] * M[2][1]
                     + M[0][1] * M[1][0] * M[2][2]
                     + M[0][2] * M[1][1] * M[2][0]))

            return sigma1, sigma2, sigma3

    # sigma1, sigma2, sigma3 represent a+b+c, ab+ac+bc, abc respectively
    # s1, s2, s3 = elem_sym_polys_of_eigen_values(C)
    # print(s1, s2, s3)

    # ortho_loss = s1 + (1 + eps) * (1 + eps) * s2 / s3 - 3 * 2 * (1 + eps)

    ortho_loss = (s1 + (1 + eps) * (1 + eps) / s1) \
               + (s2 + (1 + eps) * (1 + eps) / s2) \
               + (s3 + (1 + eps) * (1 + eps) / s3) - 2 * 3 * (1 + eps)

    return torch.mean(ortho_loss)
//  [python]
def conv_params_pdpt(ni, no, k):
    w = conv_params(ni, no, k)
    if ni != no:
        return w
    else:
        n = no
        r = int(np.ceil(float(n)/2))
        slices = w.view(n, n, -1).permute(2, 0, 1)
        es, Vs = [], []
        for u in slices:
            #u_sym = u.triu(diagonal=1).t() + u.triu()
            u_sym = (u + u.t()).div(2)
            e, V = torch.symeig(u_sym)
            es.append(e[r:-1])
            Vs.append(V[:,r:-1])
        return {'e': torch.stack(es), 'V': torch.stack(Vs)}
//  [python]
def conv_params_pdpt(ni, no, k):
    w = conv_params(ni, no, k)
    if ni != no:
        return w
    else:
        n = no
        r = int(np.ceil(float(n)/2))
        slices = w.view(n, n, -1).permute(2, 0, 1)
        es, Vs = [], []
        for u in slices:
            #u_sym = u.triu(diagonal=1).t() + u.triu()
            u_sym = (u + u.t()).div(2)
            e, V = torch.linalg.eigh(u_sym)
            es.append(e[r:-1])
            Vs.append(V[:,r:-1])
        return {'e': torch.stack(es), 'V': torch.stack(Vs)}
//  [python]
def _fit_3d(self, data: Tensor) -> None:
        _, _, columns = data.shape
        eigenvalues, eigenvectors = [], []
        for column in range(columns):
            data_column = data[:, :, column].t()
            mean = torch.mean(data_column, dim=0)
            shifted_data = data_column - mean
            cov = torch.mm(shifted_data.t(), shifted_data)
            eigenvalue, eigenvector = torch.symeig(cov, eigenvectors=True)
            eigenvalues.append(eigenvalue)
            eigenvectors.append(eigenvector)

        self.params = {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}
//  [python]
def _fit_3d(self, data: Tensor) -> None:
        _, _, columns = data.shape
        eigenvalues, eigenvectors = [], []
        for column in range(columns):
            data_column = data[:, :, column].t()
            mean = torch.mean(data_column, dim=0)
            shifted_data = data_column - mean
            cov = torch.mm(shifted_data.t(), shifted_data)
            eigenvalue, eigenvector = torch.linalg.eigh(cov)
            eigenvalues.append(eigenvalue)
            eigenvectors.append(eigenvector)

        self.params = {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}