jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, kernel, support, device='cpu'):
        assert support.shape[1] == 1, 'Only supported for d_x = 1'
        self.support = support  # (n_x, 1)
        self.device = device

        self.cov_matr = kernel(support, support).to(self.device).double()
        self.cov_matr_ddx = kernel(support, support, differentiate=True).double().to(self.device)
        dx = support[1:] - support[:-1]
        self.ddx_matrix = make_ddx_operator(support.shape[0], dx=dx).to(self.device).double()
        self.discrete_op = self.cov_matr + torch.matmul(self.cov_matr_ddx, self.ddx_matrix)
        self.inv_op = torch.linalg.inv(self.discrete_op).to(self.device)  # (n_x, n_x)
        self.sobolev_form = self.ddx_matrix.T @ self.ddx_matrix @ self.inv_op  # D^T D C^{-1}
        self.quadratic_form = self.inv_op + self.sobolev_form

        # The quadratic form is not symmetric, PSD in general. So let's project it.
        # https://nhigham.com/2021/01/26/what-is-the-nearest-positive-semidefinite-matrix/
        # https://scicomp.stackexchange.com/questions/30631/how-to-find-the-nearest-a-near-positive-definite-from-a-given-matrix
        self.symmetrized_quadratic_form = 0.5 * (self.quadratic_form + self.quadratic_form.T)
        L, V = torch.linalg.eig(self.symmetrized_quadratic_form)
        assert torch.all(torch.isreal(V)), 'Got complex eigenvectors'
        assert torch.all(torch.isreal(L)), 'Got complex eigenvalues'
        V = V.double()
        L = L.double()
        eps = 1e-5
        L[L < eps] = eps
        self.quadratic_form_proj = V @ torch.diag_embed(L) @ V.T
//  [python]
def eig(T):
    return torch.eig(T)
//  [python]
def eig(T):
    return torch.linalg.eig(T)
//  [python]
def matrix_sqrt(Sigma, inverse=False):
    eigvals, U = torch.eig(torch.inverse(Sigma) if inverse else Sigma)
    return (U@torch.diag(torch.sqrt(eigvals[:,0]))@U.T).float()
//  [python]
def matrix_sqrt(Sigma, inverse=False):
    eigvals, U = torch.linalg.eig(torch.inverse(Sigma) if inverse else Sigma)
    return (U@torch.diag(torch.sqrt(eigvals))@U.T).float()
//  [python]
def compute_hvp_accurate(self, lambda_const_param=10):
        start_time = time()
        hvp_accurate_dict={}
        for weight_name in self.val_grad_avg_dict:
            # lambda_const computation
            S=torch.zeros(len(self.tr_grad_dict.keys()))
            for tr_id in self.tr_grad_dict:
                tmp_grad = self.tr_grad_dict[tr_id][weight_name]
                S[tr_id]=torch.mean(tmp_grad**2)
            lambda_const = torch.mean(S) / lambda_const_param # layer-wise lambda

            # hvp computation (eigenvalue decomposition)
            AAt_matrix = torch.zeros(torch.outer(self.tr_grad_dict[0][weight_name].reshape(-1), 
                                                 self.tr_grad_dict[0][weight_name].reshape(-1)).shape)
            for tr_id in self.tr_grad_dict: 
                tmp_mat = torch.outer(self.tr_grad_dict[tr_id][weight_name].reshape(-1), 
                                      self.tr_grad_dict[tr_id][weight_name].reshape(-1))
                AAt_matrix += tmp_mat
                
            L, V = torch.eig(AAt_matrix)
            L, V = L.float(), V.float()
            hvp = self.val_grad_avg_dict[weight_name].reshape(-1) @ V
            hvp = (hvp / (lambda_const + L/ self.n_train)) @ V.T

            hvp_accurate_dict[weight_name] = hvp.reshape(len(self.tr_grad_dict[0][weight_name]), -1)
            del tmp_mat, AAt_matrix, V # to save memory
        self.hvp_dict['accurate'] = hvp_accurate_dict
        self.time_dict['accurate'] = time()-start_time
//  [python]
def compute_hvp_accurate(self, lambda_const_param=10):
        start_time = time()
        hvp_accurate_dict={}
        for weight_name in self.val_grad_avg_dict:
            # lambda_const computation
            S=torch.zeros(len(self.tr_grad_dict.keys()))
            for tr_id in self.tr_grad_dict:
                tmp_grad = self.tr_grad_dict[tr_id][weight_name]
                S[tr_id]=torch.mean(tmp_grad**2)
            lambda_const = torch.mean(S) / lambda_const_param # layer-wise lambda

            # hvp computation (eigenvalue decomposition)
            AAt_matrix = torch.zeros(torch.outer(self.tr_grad_dict[0][weight_name].reshape(-1), 
                                                 self.tr_grad_dict[0][weight_name].reshape(-1)).shape)
            for tr_id in self.tr_grad_dict: 
                tmp_mat = torch.outer(self.tr_grad_dict[tr_id][weight_name].reshape(-1), 
                                      self.tr_grad_dict[tr_id][weight_name].reshape(-1))
                AAt_matrix += tmp_mat
                
            L, V = torch.linalg.eig(AAt_matrix)
            L, V = L.float(), V.float()
            hvp = self.val_grad_avg_dict[weight_name].reshape(-1) @ V
            hvp = (hvp / (lambda_const + L/ self.n_train)) @ V.T

            hvp_accurate_dict[weight_name] = hvp.reshape(len(self.tr_grad_dict[0][weight_name]), -1)
            del tmp_mat, AAt_matrix, V # to save memory
        self.hvp_dict['accurate'] = hvp_accurate_dict
        self.time_dict['accurate'] = time()-start_time
//  [python]
def positional_encoding(g, pos_enc_dim):
    """
        Graph positional encoding v/ Laplacian eigenvectors
    """

    # Laplacian
    A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    A = torch.tensor(A.toarray()).to("cuda")
    N = torch.diag(g.in_degrees().clip(1) ** -0.5).double().to("cuda")
    L = torch.eye(g.number_of_nodes()).to("cuda") - N.mm(A).mm(N)

    EigVal, EigVec = torch.eig(L)
    idx = torch.real(EigVal[:,0]).argsort() # increasing order
    EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
    g.ndata['pos_enc'] = EigVec[:,1:pos_enc_dim+1].cpu().float() 
    
    # Laplacian
    # A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    # N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=float)
    # L = sp.eye(g.number_of_nodes()) - N * A * N

    # # Eigenvectors with numpy
    # EigVal, EigVec = np.linalg.eig(L.toarray())
    # idx = EigVal.argsort() # increasing order
    # EigVal, EigVec = EigVal[idx], np.real(EigVec[:,idx])
    # g.ndata['pos_enc'] = torch.from_numpy(EigVec[:,1:pos_enc_dim+1]).float() 

    return g
//  [python]
def positional_encoding(g, pos_enc_dim):
    """
        Graph positional encoding v/ Laplacian eigenvectors
    """

    # Laplacian
    A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    A = torch.tensor(A.toarray()).to("cuda")
    N = torch.diag(g.in_degrees().clip(1) ** -0.5).double().to("cuda")
    L = torch.eye(g.number_of_nodes()).to("cuda") - N.mm(A).mm(N)

    EigVal, EigVec = torch.linalg.eig(L)
    idx = torch.real(EigVal).argsort() # increasing order
    EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
    g.ndata['pos_enc'] = EigVec[:,1:pos_enc_dim+1].cpu().float() 
    
    # Laplacian
    # A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    # N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=float)
    # L = sp.eye(g.number_of_nodes()) - N * A * N

    # # Eigenvectors with numpy
    # EigVal, EigVec = np.linalg.eig(L.toarray())
    # idx = EigVal.argsort() # increasing order
    # EigVal, EigVec = EigVal[idx], np.real(EigVec[:,idx])
    # g.ndata['pos_enc'] = torch.from_numpy(EigVec[:,1:pos_enc_dim+1]).float() 

    return g
//  [python]
def PCA(self, input_tensor, PCA_dim):
        mean = torch.mean(input_tensor, dim=0)
        X_centered = input_tensor - mean.unsqueeze(0)
        X_centered = X_centered.float()
        cov_matrix = torch.mm(X_centered.T, X_centered)
        eigenvalues, eigenvectors = torch.symeig(cov_matrix, eigenvectors=True)
        eigenvalues = eigenvalues.float()
        eigenvectors = eigenvectors.float()    
        sorted_indices = torch.argsort(eigenvalues, descending=True)
        eigenvectors = eigenvectors[:, sorted_indices]
        principal_components = eigenvectors[:, :PCA_dim]
        X_transformed = torch.mm(X_centered, principal_components)
        X_reversed = torch.mm(X_transformed, principal_components.T)
        X_reversed += mean
        return X_reversed
//  [python]
def PCA(self, input_tensor, PCA_dim):
        mean = torch.mean(input_tensor, dim=0)
        X_centered = input_tensor - mean.unsqueeze(0)
        X_centered = X_centered.float()
        cov_matrix = torch.mm(X_centered.T, X_centered)
        eigenvalues, eigenvectors = torch.linalg.eig(cov_matrix)
        eigenvalues = eigenvalues.float()
        eigenvectors = eigenvectors.float()    
        sorted_indices = torch.argsort(eigenvalues, descending=True)
        eigenvectors = eigenvectors[:, sorted_indices]
        principal_components = eigenvectors[:, :PCA_dim]
        X_transformed = torch.mm(X_centered, principal_components)
        X_reversed = torch.mm(X_transformed, principal_components.T)
        X_reversed += mean
        return X_reversed
//  [python]
def compute_auxillary_matrices(self):
        # A operator
        H, g_prime = self.H.squeeze(), self.g_prime.squeeze()
        self.A = torch.zeros((self.nl,self.nl), **self.arr_kwargs)
        if self.nl == 1:
            self.A[0,0] = 1./(H*g_prime)
        else:
            self.A[0,0] = 1./(H[0]*g_prime[0]) + 1./(H[0]*g_prime[1])
            self.A[0,1] = -1./(H[0]*g_prime[1])
            for i in range(1, self.nl-1):
                self.A[i,i-1] = -1./(H[i]*g_prime[i])
                self.A[i,i] = 1./H[i]*(1/g_prime[i+1] + 1/g_prime[i])
                self.A[i,i+1] = -1./(H[i]*g_prime[i+1])
            self.A[-1,-1] = 1./(H[self.nl-1]*g_prime[self.nl-1])
            self.A[-1,-2] = -1./(H[self.nl-1]*g_prime[self.nl-1])

        # layer-to-mode and mode-to-layer matrices
        ev_A, P = torch.eig(self.A)
        self.lambda_sq = self.f0**2 * ev_A[:,0].real.reshape((1, self.nl, 1, 1))
        self.Cl2m = torch.inverse(P.real)
        self.Cm2l = P.real
        self.rossby_radii = self.lambda_sq.squeeze().pow(-0.5)
        with np.printoptions(precision=1):
            print(f'Rossby rad.: {self.rossby_radii.cpu().numpy()/1e3} km')

        # For Helmholtz equations
        nx, ny, nl = self.nx, self.ny, self.nl
        laplace_dst = compute_laplace_dst(
                nx, ny, self.dx, self.dy, self.arr_kwargs).unsqueeze(0).unsqueeze(0)
        self.helmholtz_dst =  laplace_dst - self.lambda_sq


        # homogeneous Helmholtz solutions
        cst = torch.ones((1, nl, nx+1, ny+1), **self.arr_kwargs)
        if len(self.masks.psi_irrbound_xids) > 0:
            self.cap_matrices = compute_capacitance_matrices(
                self.helmholtz_dst, self.masks.psi_irrbound_xids,
                self.masks.psi_irrbound_yids)
            sol = solve_helmholtz_dst_cmm(
                    (cst*self.masks.psi)[...,1:-1,1:-1],
                    self.helmholtz_dst, self.cap_matrices,
                    self.masks.psi_irrbound_xids,
                    self.masks.psi_irrbound_yids,
                    self.masks.psi)
        else:
            self.cap_matrices = None
            sol = solve_helmholtz_dst(cst[...,1:-1,1:-1], self.helmholtz_dst)

        self.homsol = cst + sol * self.lambda_sq
        self.homsol_mean = (
                interp_TP(self.homsol)*self.masks.q).mean((-1,-2), keepdim=True)
        self.helmholtz_dst = self.helmholtz_dst.type(torch.float32)
//  [python]
def compute_auxillary_matrices(self):
        # A operator
        H, g_prime = self.H.squeeze(), self.g_prime.squeeze()
        self.A = torch.zeros((self.nl,self.nl), **self.arr_kwargs)
        if self.nl == 1:
            self.A[0,0] = 1./(H*g_prime)
        else:
            self.A[0,0] = 1./(H[0]*g_prime[0]) + 1./(H[0]*g_prime[1])
            self.A[0,1] = -1./(H[0]*g_prime[1])
            for i in range(1, self.nl-1):
                self.A[i,i-1] = -1./(H[i]*g_prime[i])
                self.A[i,i] = 1./H[i]*(1/g_prime[i+1] + 1/g_prime[i])
                self.A[i,i+1] = -1./(H[i]*g_prime[i+1])
            self.A[-1,-1] = 1./(H[self.nl-1]*g_prime[self.nl-1])
            self.A[-1,-2] = -1./(H[self.nl-1]*g_prime[self.nl-1])

        # layer-to-mode and mode-to-layer matrices
        ev_A, P = torch.linalg.eig(self.A)
        self.lambda_sq = self.f0**2 * ev_A.real.reshape((1, self.nl, 1, 1))
        self.Cl2m = torch.linalg.inv(P.real)
        self.Cm2l = P.real
        self.rossby_radii = self.lambda_sq.squeeze().pow(-0.5)
        with np.printoptions(precision=1):
            print(f'Rossby rad.: {self.rossby_radii.cpu().numpy()/1e3} km')

        # For Helmholtz equations
        nx, ny, nl = self.nx, self.ny, self.nl
        laplace_dst = compute_laplace_dst(
                nx, ny, self.dx, self.dy, self.arr_kwargs).unsqueeze(0).unsqueeze(0)
        self.helmholtz_dst =  laplace_dst - self.lambda_sq


        # homogeneous Helmholtz solutions
        cst = torch.ones((1, nl, nx+1, ny+1), **self.arr_kwargs)
        if len(self.masks.psi_irrbound_xids) > 0:
            self.cap_matrices = compute_capacitance_matrices(
                self.helmholtz_dst, self.masks.psi_irrbound_xids,
                self.masks.psi_irrbound_yids)
            sol = solve_helmholtz_dst_cmm(
                    (cst*self.masks.psi)[...,1:-1,1:-1],
                    self.helmholtz_dst, self.cap_matrices,
                    self.masks.psi_irrbound_xids,
                    self.masks.psi_irrbound_yids,
                    self.masks.psi)
        else:
            self.cap_matrices = None
            sol = solve_helmholtz_dst(cst[...,1:-1,1:-1], self.helmholtz_dst)

        self.homsol = cst + sol * self.lambda_sq
        self.homsol_mean = (
                interp_TP(self.homsol)*self.masks.q).mean((-1,-2), keepdim=True)
        self.helmholtz_dst = self.helmholtz_dst.type(torch.float32)
//  [python]
def matrix_operations(A):
    """Calculate the -1/2 power of matrix A"""

    V, Q = t.eig(A)
    V_inverse = diag(V ** (-0.5))
    A_inverse = mm(mm(Q, V_inverse), t.inverse(Q))

    return A_inverse.double()
//  [python]
def matrix_operations(A):
    """Calculate the -1/2 power of matrix A"""

    V, Q = linalg.eig(A)
    V_inverse = diag(V ** (-0.5))
    A_inverse = mm(mm(Q, V_inverse), linalg.inv(Q))

    return A_inverse.double()
//  [python]
def laplacian_positional_enc(self, adj):
        A = adj
        D = torch.diag(torch.count_nonzero(A, dim=-1))
        L = torch.eye(A.shape[0], device=A.device) - D * A * D

        EigVal, EigVec = torch.eig(L, eigenvectors=True)
        idx = torch.argsort(EigVal[:,0])
        EigVal, EigVec = EigVal[idx], EigVec[:,idx]
        pos_enc = EigVec[:,1:self.pos_enc_dim + 1]
        return pos_enc
//  [python]
def laplacian_positional_enc(self, adj):
        A = adj
        D = torch.diag(torch.count_nonzero(A, dim=-1))
        L = torch.eye(A.shape[0], device=A.device) - D * A * D

        EigVal, EigVec = torch.linalg.eig(L)
        idx = torch.argsort(torch.real(EigVal))
        EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
        pos_enc = EigVec[:,1:self.pos_enc_dim + 1]
        return pos_enc
//  [python]
def logm(A):
    """Calculate the matrix logarithm of matrix A"""

    V, Q = t.eig(A)  # V为特征值,Q为特征向量
    V_log = diag(log(V))
    A_logm = mm(mm(Q, V_log), linalg.inv(Q))

    return A_logm.double()
//  [python]
def logm(A):
    """Calculate the matrix logarithm of matrix A"""

    V, Q = linalg.eig(A)  # V为特征值,Q为特征向量
    V_log = diag(log(V))
    A_logm = mm(mm(Q, V_log), linalg.inv(Q))

    return A_logm.double()
//  [python]
def get_log_like_gradient_and_hessian(
        self, X, T, P, S, Y_target, X_min, X_max, eps=1e-2, compute_hessian=False
    ):

        X_bar = (X - X_min) / (X_max - X_min)
        log_pi = self.V(X, T, P, S, Y_target, X_bar)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.stack(
                [torch.autograd.grad(e, X, retain_graph=True)[0] for e in g]
            )
            lamda, Q = torch.eig(H, eigenvectors=True)
            lamda, Q = lamda.type(torch.float), Q.type(torch.float)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / torch.sqrt(lamda**2 + eps)
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def get_log_like_gradient_and_hessian(
        self, X, T, P, S, Y_target, X_min, X_max, eps=1e-2, compute_hessian=False
    ):

        X_bar = (X - X_min) / (X_max - X_min)
        log_pi = self.V(X, T, P, S, Y_target, X_bar)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.stack(
                [torch.autograd.grad(e, X, retain_graph=True)[0] for e in g]
            )
            lamda, Q = torch.linalg.eig(H)
            lamda, Q = lamda.type(torch.float), Q.type(torch.float)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / torch.sqrt(lamda**2 + eps)
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def translation_from_gram_matrix(gram_matrices: torch.Tensor) -> torch.Tensor:
    eigenvectors = []
    if gram_matrices.ndim == 2:
        gram_matrices = gram_matrices[None]
    for gram_matrix in gram_matrices:
        # TODO: with amp error
        L, V = torch.eig(gram_matrix)
        V = V.real
        real_eigenvalues = L[:,0]
        min_eigenvalue = real_eigenvalues[0]
        min_eigenvectors = V[:, 0] / torch.linalg.norm(V[:, 0])
        for idx, eigenvalue in enumerate(real_eigenvalues):
            if eigenvalue < min_eigenvalue:
                min_eigenvalue = eigenvalue
                min_eigenvectors = V[:, idx] / torch.linalg.norm(V[:, idx])
        eigenvectors.append(min_eigenvectors)
    return torch.stack(eigenvectors, dim=0)
//  [python]
def translation_from_gram_matrix(gram_matrices: torch.Tensor) -> torch.Tensor:
    eigenvectors = []
    if gram_matrices.ndim == 2:
        gram_matrices = gram_matrices[None]
    for gram_matrix in gram_matrices:
        # TODO: with amp error
        L, V = torch.linalg.eig(gram_matrix)
        V = V.real
        real_eigenvalues = L.real
        min_eigenvalue = real_eigenvalues[0]
        min_eigenvectors = V[:, 0] / torch.linalg.norm(V[:, 0])
        for idx, eigenvalue in enumerate(real_eigenvalues):
            if eigenvalue < min_eigenvalue:
                min_eigenvalue = eigenvalue
                min_eigenvectors = V[:, idx] / torch.linalg.norm(V[:, idx])
        eigenvectors.append(min_eigenvectors)
    return torch.stack(eigenvectors, dim=0)
//  [python]
def get_log_like_gradient_and_hessian(
        self, X, X_I, Y_target, X_min, X_max, eps=1e-2, compute_hessian=False
    ):

        X_bar = (X - X_min) / (X_max - X_min)
        log_pi = self.V(X, X_I, Y_target, X_bar)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.stack(
                [torch.autograd.grad(e, X, retain_graph=True)[0] for e in g]
            )
            lamda, Q = torch.eig(H, eigenvectors=True)
            lamda, Q = lamda.type(torch.float), Q.type(torch.float)
            lamda_prime = torch.sqrt(lamda[:,0]**2 + eps)
            lamda_prime_inv = 1.0 / torch.sqrt(lamda[:,0]**2 + eps)
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def get_log_like_gradient_and_hessian(
        self, X, X_I, Y_target, X_min, X_max, eps=1e-2, compute_hessian=False
    ):

        X_bar = (X - X_min) / (X_max - X_min)
        log_pi = self.V(X, X_I, Y_target, X_bar)
        if compute_hessian:
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.stack(
                [torch.autograd.grad(e, X, retain_graph=True)[0] for e in g]
            )
            lamda, Q = torch.linalg.eig(H)
            lamda, Q = lamda.type(torch.float), Q.type(torch.float)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / torch.sqrt(lamda**2 + eps)
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def _complex_python_udf(t1, t2):
        t3 = torch.nn.functional.linear(t1, t2)
        t4 = torch.nn.functional.linear(t2, t3)
        t5 = torch.nn.functional.linear(t3, t4)
        return torch.chain_matmul(t1, t2, t3, t4, t5)
//  [python]
def _complex_python_udf(t1, t2):
        t3 = torch.nn.functional.linear(t1, t2)
        t4 = torch.nn.functional.linear(t2, t3)
        t5 = torch.nn.functional.linear(t3, t4)
        return torch.linalg.multi_dot([t1, t2, t3, t4, t5])
//  [python]
def torch_pca_eig(X, center=True, scale=False, zero_negatives=False):
    """
    perform PCA over X
    - X is of size (num_points, vec_size)

    NOTE: unlike PCA_svd, the number of vectors/values returned is always: vec_size

    WARNING: this may be incorrect!
    """
    n, m = X.shape
    # center points along axes
    if center:
        X = X - X.mean(dim=0)
    # compute covariance -- TODO: optimise this line
    covariance = (1 / (n - 1)) * torch.mm(X.T, X)
    if scale:
        scaling = torch.sqrt(1 / torch.diagonal(covariance))
        covariance = torch.mm(torch.diagflat(scaling), covariance)
    # compute eigen values and eigen vectors
    eigenvalues, eigenvectors = torch.symeig(covariance, eigenvectors=True)
    # sort components by decreasing variance
    components = torch.real(eigenvectors.T)  # TODO: handle imaginary numbers!
    explained_variance = torch.real(eigenvalues)  # TODO: handle imaginary numbers!
    # handle n < m -- numerical stability issues return negative values!
    #                 maybe this should just zero out the negatives instead, they don't contribute!?
    explained_variance = torch.abs(explained_variance)
    # return sorted
    idxs = torch.argsort(explained_variance, descending=True)
    return components[idxs], explained_variance[idxs]
//  [python]
def torch_pca_eig(X, center=True, scale=False, zero_negatives=False):
    """
    perform PCA over X
    - X is of size (num_points, vec_size)

    NOTE: unlike PCA_svd, the number of vectors/values returned is always: vec_size

    WARNING: this may be incorrect!
    """
    n, m = X.shape
    # center points along axes
    if center:
        X = X - X.mean(dim=0)
    # compute covariance -- TODO: optimise this line
    covariance = (1 / (n - 1)) * torch.mm(X.T, X)
    if scale:
        scaling = torch.sqrt(1 / torch.diagonal(covariance))
        covariance = torch.mm(torch.diagflat(scaling), covariance)
    # compute eigen values and eigen vectors
    eigenvalues, eigenvectors = torch.linalg.eig(covariance)
    # sort components by decreasing variance
    components = torch.real(eigenvectors.T)  # TODO: handle imaginary numbers!
    explained_variance = torch.real(eigenvalues)  # TODO: handle imaginary numbers!
    # handle n < m -- numerical stability issues return negative values!
    #                 maybe this should just zero out the negatives instead, they don't contribute!?
    explained_variance = torch.abs(explained_variance)
    # return sorted
    idxs = torch.argsort(explained_variance, descending=True)
    return components[idxs], explained_variance[idxs]
//  [python]
def pca(x, dim=2):
    """
    Perform PCA to for dimensionality reduction
    
    Arguments:
    ----------
    x: torch.Tensor
        Input data of size (N, K)
    dim: int
        Desired output dimension
        
    Returns:
    --------
    torch.Tensor
        Output data of size (N, dim)
    torch.Tensor
        Selected principal components of size (K, dim)
    """
    n = len(x)
    x_ = x - x.mean(0, keepdim=True)
    gram = x_.matmul(x_.T) / n
    e, v = torch.eig(gram, eigenvectors=True)
    e = torch.real(e[:,0])
    v = torch.real(v)
    eigvec = (x_.T).matmul(v) / torch.sqrt(n * e.unsqueeze(0))

    order = torch.argsort(e, descending=True)
    pc = eigvec[:, order[:dim]]

    return x_.matmul(pc), pc
//  [python]
def pca(x, dim=2):
    """
    Perform PCA to for dimensionality reduction
    
    Arguments:
    ----------
    x: torch.Tensor
        Input data of size (N, K)
    dim: int
        Desired output dimension
        
    Returns:
    --------
    torch.Tensor
        Output data of size (N, dim)
    torch.Tensor
        Selected principal components of size (K, dim)
    """
    n = len(x)
    x_ = x - x.mean(0, keepdim=True)
    gram = x_.matmul(x_.T) / n
    e, v = torch.linalg.eig(gram)
    e = torch.real(e)
    v = torch.real(v)
    eigvec = (x_.T).matmul(v) / torch.sqrt(n * e.unsqueeze(0))

    order = torch.argsort(e, descending=True)
    pc = eigvec[:, order[:dim]]

    return x_.matmul(pc), pc
//  [python]
def _eig_helper(H: Tensor):
    """
    wrapper around `torch.linalg.eig` that sorts eigenvalues / eigenvectors by
    ascending eigenvalues, like `torch.linalg.eigh`, and returns the real component
    (since `H` is never complex, there should never be a complex component. however,
    `torch.linalg.eig` always returns a complex tensor, which in this case would
    actually have no complex component)
    """
    ls, vs = torch.eig(H)
    ls, vs = ls.real, vs.real
    ls_argsort = torch.argsort(ls)
    vs = vs[:, ls_argsort]
    ls = ls[ls_argsort]
    return ls, vs
//  [python]
def _eig_helper(H: Tensor):
    """
    wrapper around `torch.linalg.eig` that sorts eigenvalues / eigenvectors by
    ascending eigenvalues, like `torch.linalg.eigh`, and returns the real component
    (since `H` is never complex, there should never be a complex component. however,
    `torch.linalg.eig` always returns a complex tensor, which in this case would
    actually have no complex component)
    """
    ls, vs = torch.linalg.eig(H)
    ls, vs = ls.real, vs.real
    ls_argsort = torch.argsort(ls)
    vs = vs[:, ls_argsort]
    ls = ls[ls_argsort]
    return ls, vs
//  [python]
def calculate_fix(frame1, frame2):
    mat1_0, mat1_1 = torch.eye(4), torch.eye(4)
    mat2_0, mat2_1 = torch.eye(4), torch.eye(4)
    mat1_0[:3, :3], mat1_0[:3, 3] = torch.from_numpy(R.from_rotvec(frame1[0, 3:]).as_matrix()), frame1[0, :3]
    mat1_1[:3, :3], mat1_1[:3, 3] = torch.from_numpy(R.from_rotvec(frame1[1, 3:]).as_matrix()), frame1[1, :3]
    mat2_0[:3, :3], mat2_0[:3, 3] = torch.from_numpy(R.from_rotvec(frame2[0, 3:]).as_matrix()), frame2[0, :3]
    mat2_1[:3, :3], mat2_1[:3, 3] = torch.from_numpy(R.from_rotvec(frame2[1, 3:]).as_matrix()), frame2[1, :3]

    mat1 = torch.matmul(torch.linalg.inv(mat1_0), mat1_1)
    mat2 = torch.matmul(torch.linalg.inv(mat2_0), mat2_1)

    d_mat = mat2.matmul(torch.inverse(mat1))
    L, V = torch.eig(d_mat)
    fix = None
    for i in range(4):
        if abs(L[i,0]-1)<1e-5 and abs(V[3, i])>1e-5:
            fix = (V[:3, i] / V[3, i]).real
    if fix is None:
        print('part interpolate error')

    rot_vec = torch.tensor(R.from_matrix(d_mat[:3, :3]).as_rotvec()).float()
    return fix, rot_vec
//  [python]
def calculate_fix(frame1, frame2):
    mat1_0, mat1_1 = torch.eye(4), torch.eye(4)
    mat2_0, mat2_1 = torch.eye(4), torch.eye(4)
    mat1_0[:3, :3], mat1_0[:3, 3] = torch.from_numpy(R.from_rotvec(frame1[0, 3:]).as_matrix()), frame1[0, :3]
    mat1_1[:3, :3], mat1_1[:3, 3] = torch.from_numpy(R.from_rotvec(frame1[1, 3:]).as_matrix()), frame1[1, :3]
    mat2_0[:3, :3], mat2_0[:3, 3] = torch.from_numpy(R.from_rotvec(frame2[0, 3:]).as_matrix()), frame2[0, :3]
    mat2_1[:3, :3], mat2_1[:3, 3] = torch.from_numpy(R.from_rotvec(frame2[1, 3:]).as_matrix()), frame2[1, :3]

    mat1 = torch.matmul(torch.linalg.inv(mat1_0), mat1_1)
    mat2 = torch.matmul(torch.linalg.inv(mat2_0), mat2_1)

    d_mat = mat2.matmul(torch.inverse(mat1))
    L, V = torch.linalg.eig(d_mat)
    fix = None
    for i in range(4):
        if abs(L[i]-1)<1e-5 and abs(V[3, i])>1e-5:
            fix = (V[:3, i] / V[3, i]).real
    if fix is None:
        print('part interpolate error')

    rot_vec = torch.tensor(R.from_matrix(d_mat[:3, :3]).as_rotvec()).float()
    return fix, rot_vec
//  [python]
def log_histograms(self, trainer, epoch_num, representations):
        tensorboard = trainer.logger.experiment

        X = representations
        num_units = X.shape[-1]
        num_samples = X.shape[0]
        X_std = (X - X.mean(dim=0, keepdim=True)) / (X.std(dim=0, keepdim=True) + 1e-8)

        mean_sq_activity = (X ** 2).mean(dim=0)

        corr = torch.matmul(X_std.T, X_std) / (num_samples - 1)
        off_diag_idxs = torch.triu_indices(num_units, num_units, offset=1)
        cross_unit_corrs = corr[off_diag_idxs[0], off_diag_idxs[1]]

        eig_vals_corr, _ = torch.eig(corr)
        eig_vals_corr = eig_vals_corr[:,0]
        eig_vals_corr_norm = eig_vals_corr.clamp(1e-8) / eig_vals_corr.max()

        gram = torch.matmul(X.T, X)
        eig_vals_gram, _ = torch.eig(gram)
        eig_vals_gram = eig_vals_gram[:,0]
        eig_vals_gram_norm = eig_vals_gram.clamp(1e-8) / eig_vals_gram.max()

        cos_dists = pairwise_cosine_similarity(X - X.mean(dim=0))
        off_diag_idxs = torch.triu_indices(num_samples, num_samples, offset=1)
        cos_dists = cos_dists[off_diag_idxs[0], off_diag_idxs[1]]

        tensorboard.add_histogram('Activity Statistics/unit_activities', mean_sq_activity,
                                  global_step=epoch_num)
        tensorboard.add_histogram('Activity Statistics/unit_correlations', cross_unit_corrs,
                                  global_step=epoch_num)

        tensorboard.add_histogram('Correlation Matrix/eig_vals', eig_vals_corr, global_step=epoch_num)
        tensorboard.add_histogram('Correlation Matrix/normalized_eig_vals', torch.log(eig_vals_corr_norm),
                                  global_step=epoch_num)

        tensorboard.add_histogram('Gramian Matrix/eig_vals', eig_vals_gram, global_step=epoch_num)
        tensorboard.add_histogram('Gramian Matrix/normalized_eig_vals', torch.log(eig_vals_gram_norm),
                                  global_step=epoch_num)

        tensorboard.add_histogram('Embedding Similarity/cos_dists', cos_dists, global_step=epoch_num)
//  [python]
def log_histograms(self, trainer, epoch_num, representations):
        tensorboard = trainer.logger.experiment

        X = representations
        num_units = X.shape[-1]
        num_samples = X.shape[0]
        X_std = (X - X.mean(dim=0, keepdim=True)) / (X.std(dim=0, keepdim=True) + 1e-8)

        mean_sq_activity = (X ** 2).mean(dim=0)

        corr = torch.matmul(X_std.T, X_std) / (num_samples - 1)
        off_diag_idxs = torch.triu_indices(num_units, num_units, offset=1)
        cross_unit_corrs = corr[off_diag_idxs[0], off_diag_idxs[1]]

        eig_vals_corr, _ = torch.linalg.eig(corr)
        eig_vals_corr = eig_vals_corr.real
        eig_vals_corr_norm = eig_vals_corr.clamp(1e-8) / eig_vals_corr.max()

        gram = torch.matmul(X.T, X)
        eig_vals_gram, _ = torch.linalg.eig(gram)
        eig_vals_gram = eig_vals_gram.real
        eig_vals_gram_norm = eig_vals_gram.clamp(1e-8) / eig_vals_gram.max()

        cos_dists = pairwise_cosine_similarity(X - X.mean(dim=0))
        off_diag_idxs = torch.triu_indices(num_samples, num_samples, offset=1)
        cos_dists = cos_dists[off_diag_idxs[0], off_diag_idxs[1]]

        tensorboard.add_histogram('Activity Statistics/unit_activities', mean_sq_activity,
                                  global_step=epoch_num)
        tensorboard.add_histogram('Activity Statistics/unit_correlations', cross_unit_corrs,
                                  global_step=epoch_num)

        tensorboard.add_histogram('Correlation Matrix/eig_vals', eig_vals_corr, global_step=epoch_num)
        tensorboard.add_histogram('Correlation Matrix/normalized_eig_vals', torch.log(eig_vals_corr_norm),
                                  global_step=epoch_num)

        tensorboard.add_histogram('Gramian Matrix/eig_vals', eig_vals_gram, global_step=epoch_num)
        tensorboard.add_histogram('Gramian Matrix/normalized_eig_vals', torch.log(eig_vals_gram_norm),
                                  global_step=epoch_num)

        tensorboard.add_histogram('Embedding Similarity/cos_dists', cos_dists, global_step=epoch_num)
//  [python]
def _python_udf_with_backward_error(t1, t2):
        t3 = t1 + t2
        t4 = SimulateBackwardError.apply(t3)
        return torch.chain_matmul(t1, t2, t3, t4)
//  [python]
def _python_udf_with_backward_error(t1, t2):
        t3 = t1 + t2
        t4 = SimulateBackwardError.apply(t3)
        return torch.linalg.multi_dot([t1, t2, t3, t4])
//  [python]
def _corr_mat(self, corr, num_zern=36):
        '''
        This function generates the correlation matrix for zernike coefficients associated with
            higher-order abberations for the entire image
        Input: 
            corr       -  Correlation strength. suggested range: (-1 ~ -0.01), with -1 has the
                        weakest correlation and -0.01 has the strongest.
            num_zern   -  number of zernike coefficients. Default: 36
        '''
        # subC = _nollCovMat(num_zern, self.D, self.r0).cuda()
        subC = _nollCovMat(num_zern, 1, 1).cuda()  # current setting
        e_val, e_vec = torch.eig(subC)
        R_z = torch.real(e_vec * torch.sqrt(e_val))
        h = torch.linspace(0, self.gh-1, self.gh).cuda() - self.gh/2
        w = torch.linspace(0, self.gw-1, self.gw).cuda() - self.gw/2
        yv, xv = torch.meshgrid(h, w)
        # dist = torch.exp(corr*(xv.abs()+yv.abs()))
        dist = torch.exp(corr*(xv**2 + yv**2)**0.5)
        psd = torch.fft.fft2(dist)
        return R_z, torch.sqrt(psd)
//  [python]
def _corr_mat(self, corr, num_zern=36):
        '''
        This function generates the correlation matrix for zernike coefficients associated with
            higher-order abberations for the entire image
        Input: 
            corr       -  Correlation strength. suggested range: (-1 ~ -0.01), with -1 has the
                        weakest correlation and -0.01 has the strongest.
            num_zern   -  number of zernike coefficients. Default: 36
        '''
        # subC = _nollCovMat(num_zern, self.D, self.r0).cuda()
        subC = _nollCovMat(num_zern, 1, 1).cuda()  # current setting
        e_val, e_vec = torch.linalg.eig(subC)
        R_z = torch.real(e_vec * torch.sqrt(e_val))
        h = torch.linspace(0, self.gh-1, self.gh).cuda() - self.gh/2
        w = torch.linspace(0, self.gw-1, self.gw).cuda() - self.gw/2
        yv, xv = torch.meshgrid(h, w)
        # dist = torch.exp(corr*(xv.abs()+yv.abs()))
        dist = torch.exp(corr*(xv**2 + yv**2)**0.5)
        psd = torch.fft.fft2(dist)
        return R_z, torch.sqrt(psd)
//  [python]
def _nested_python_udf(t1, t2, dst):
        t3 = t1 * t2
        t4 = t1 + t2
        res = rpc.rpc_sync(worker_name(dst), my_py_add, args=(t3, t4))
        return torch.chain_matmul(t1, t2, t3, t4, res)
//  [python]
def _nested_python_udf(t1, t2, dst):
        t3 = t1 * t2
        t4 = t1 + t2
        res = rpc.rpc_sync(worker_name(dst), my_py_add, args=(t3, t4))
        return torch.linalg.multi_dot([t1, t2, t3, t4, res])
//  [python]
def _nested_rpc_call_backward_error(t1, t2, dst):
        t1 = t1 * t2
        t2 = t1 + t2
        res = rpc.rpc_sync(
            worker_name(dst),
            DistAutogradTest._python_udf_with_backward_error,
            args=(t1, t2),
        )
        return torch.chain_matmul(t1, t2, res)
//  [python]
def _nested_rpc_call_backward_error(t1, t2, dst):
        t1 = t1 * t2
        t2 = t1 + t2
        res = rpc.rpc_sync(
            worker_name(dst),
            DistAutogradTest._python_udf_with_backward_error,
            args=(t1, t2),
        )
        return torch.linalg.multi_dot([t1, t2, res])
//  [python]
def test_backwards_nested_python_udf(self):
        # Run equivalent of _nested_python_udf locally.
        t1 = torch.rand((3, 3), requires_grad=True)
        t2 = torch.rand((3, 3), requires_grad=True)
        t3 = t1 * t2
        t4 = t1 + t2
        res = t3 + t4
        loss = torch.chain_matmul(t1, t2, t3, t4, res).sum()
        torch.autograd.backward([loss])

        # Now run distributed autograd.
        with dist_autograd.context() as context_id:
            loss = rpc.rpc_sync(
                worker_name(self._next_rank()),
                DistAutogradTest._nested_python_udf,
                args=(t1, t2, self._next_rank()),
            )
            dist_autograd.backward(context_id, [loss.sum()])

            grads = dist_autograd.get_gradients(context_id)
            self.assertEqual(t1.grad, grads[t1])
            self.assertEqual(t2.grad, grads[t2])
//  [python]
def test_backwards_nested_python_udf(self):
        # Run equivalent of _nested_python_udf locally.
        t1 = torch.rand((3, 3), requires_grad=True)
        t2 = torch.rand((3, 3), requires_grad=True)
        t3 = t1 * t2
        t4 = t1 + t2
        res = t3 + t4
        loss = torch.linalg.multi_dot([t1, t2, t3, t4, res]).sum()
        torch.autograd.backward([loss])

        # Now run distributed autograd.
        with dist_autograd.context() as context_id:
            loss = rpc.rpc_sync(
                worker_name(self._next_rank()),
                DistAutogradTest._nested_python_udf,
                args=(t1, t2, self._next_rank()),
            )
            dist_autograd.backward(context_id, [loss.sum()])

            grads = dist_autograd.get_gradients(context_id)
            self.assertEqual(t1.grad, grads[t1])
            self.assertEqual(t2.grad, grads[t2])
//  [python]
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        # Tries to returns a vector `v` s.t. `u = normalize(W @ v)`
        # (the invariant at top of this class) and `u @ W @ v = sigma`.
        # This uses pinverse in case W^T W is not invertible.
        v = torch.chain_matmul(weight_mat.t().mm(weight_mat).pinverse(),
                               weight_mat.t(), u.unsqueeze(1)).squeeze(1)
        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))
//  [python]
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        v = torch.linalg.multi_dot([weight_mat.t().mm(weight_mat).pinverse(),
                               weight_mat.t(), u.unsqueeze(1)]).squeeze(1)
        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))
//  [python]
def pca_torch(input_array: np.ndarray, n_components: int, test_array=None, return_eigenvectors=False,
              reg_w=0.0, device='cuda', print_flag=False):
  # 将输入的numpy数组转换为PyTorch张量
  input_array = input_array.T
  tensor = torch.from_numpy(input_array).to(device=device)
  input_mean = tensor.mean(dim=0)
  # 将输入张量中心化
  tensor = tensor - input_mean
  # 计算协方差矩阵
  cov_matrix = torch.mm(tensor.T, tensor) / (tensor.shape[0] - 1)
  cov_matrix = (cov_matrix + cov_matrix.T) / 2
  # 对协方差矩阵进行特征值分解
  eigenvalues, eigenvectors = torch.symeig(cov_matrix, eigenvectors=True)
  # Sort the eigenvalues in descending order and get the sorted indices
  sorted_indices = torch.argsort(eigenvalues, descending=True)
  # Sort the eigenvectors using the sorted indices
  eigenvectors = eigenvectors[:, sorted_indices]
  if isinstance(n_components, int):
    # 获取前n个主成分
    eigenvectors = eigenvectors[:, :n_components]
    # 计算主成分能量占比
    energy_ratio = eigenvalues[:n_components].sum() / eigenvalues.sum()
    if print_flag:
      print(energy_ratio)
  else:
    assert 0 <= n_components <= 1
    energy_ratio_cum = 0
    eigenvalues_sum = eigenvalues.sum()
    for i, e in enumerate(eigenvalues):
      energy_ratio_cum += e / eigenvalues_sum
      if energy_ratio_cum > n_components:
        n_components = i + 1
        break
    if print_flag:
      print(n_components)
    eigenvectors = eigenvectors[:, :n_components]
  # 将数据投影到主成分上
  result = torch.mm(tensor, eigenvectors)
  if test_array is not None:
    test_array = torch.from_numpy(test_array.T).to(device=device)
    test_array -= input_mean
    test_result = torch.mm(test_array, eigenvectors)
    if not return_eigenvectors:
      return result.cpu().numpy().T, test_result.cpu().numpy().T
    else:
      return result.cpu().numpy().T, test_result.cpu().numpy().T, (input_mean.cpu().numpy()[:, np.newaxis], eigenvectors.cpu().numpy().T)
  else:
    if not return_eigenvectors:
      return result.cpu().numpy().T
    else:
      return result.cpu().numpy().T, (input_mean.cpu().numpy()[:, np.newaxis], eigenvectors.cpu().numpy().T)
//  [python]
def pca_torch(input_array: np.ndarray, n_components: int, test_array=None, return_eigenvectors=False,
              reg_w=0.0, device='cuda', print_flag=False):
  # 将输入的numpy数组转换为PyTorch张量
  input_array = input_array.T
  tensor = torch.from_numpy(input_array).to(device=device)
  input_mean = tensor.mean(dim=0)
  # 将输入张量中心化
  tensor = tensor - input_mean
  # 计算协方差矩阵
  cov_matrix = torch.mm(tensor.T, tensor) / (tensor.shape[0] - 1)
  cov_matrix = (cov_matrix + cov_matrix.T) / 2
  # 对协方差矩阵进行特征值分解
  eigenvalues, eigenvectors = torch.linalg.eig(cov_matrix)
  eigenvalues = torch.real(eigenvalues)
  eigenvectors = torch.real(eigenvectors)
  # Sort the eigenvalues in descending order and get the sorted indices
  sorted_indices = torch.argsort(eigenvalues, descending=True)
  # Sort the eigenvectors using the sorted indices
  eigenvectors = eigenvectors[:, sorted_indices]
  if isinstance(n_components, int):
    # 获取前n个主成分
    eigenvectors = eigenvectors[:, :n_components]
    # 计算主成分能量占比
    eigenvalues = torch.real(eigenvalues)
    energy_ratio = eigenvalues[:n_components].sum() / eigenvalues.sum()
    if print_flag:
      print(energy_ratio)
  else:
    assert 0 <= n_components <= 1
    energy_ratio_cum = 0
    eigenvalues_sum = eigenvalues.sum()
    for i, e in enumerate(eigenvalues):
      energy_ratio_cum += e / eigenvalues_sum
      if energy_ratio_cum > n_components:
        n_components = i + 1
        break
    if print_flag:
      print(n_components)
    eigenvectors = eigenvectors[:, :n_components]
  # 将数据投影到主成分上
  result = torch.mm(tensor, eigenvectors)
  if test_array is not None:
    test_array = torch.from_numpy(test_array.T).to(device=device)
    test_array -= input_mean
    test_result = torch.mm(test_array, eigenvectors)
    if not return_eigenvectors:
      return result.cpu().numpy().T, test_result.cpu().numpy().T
    else:
      return result.cpu().numpy().T, test_result.cpu().numpy().T, (input_mean.cpu().numpy()[:, np.newaxis], eigenvectors.cpu().numpy().T)
  else:
    if not return_eigenvectors:
      return result.cpu().numpy().T
    else:
      return result.cpu().numpy().T, (input_mean.cpu().numpy()[:, np.newaxis], eigenvectors.cpu().numpy().T)
//  [python]
def _get_weights(self):
        V = self.weights
        V_t = self.weights.transpose(1, 2)
        U = self.I - 2 * torch.bmm(V, V_t) / torch.bmm(V_t, V)
        W = torch.chain_matmul(*U)

        return W
//  [python]
def _get_weights(self):
        V = self.weights
        V_t = self.weights.transpose(1, 2)
        U = self.I - 2 * torch.bmm(V, V_t) / torch.bmm(V_t, V)
        W = torch.linalg.multi_dot(U)

        return W
//  [python]
def build_int_matrix(horizon, diagonal=0, device='cpu', dtype=torch.float32, order=1,
                     traj_dt=None):
    integrate_matrix = torch.tril(torch.ones((horizon, horizon), device=device, dtype=dtype), diagonal=diagonal)
    chain_list = [torch.eye(horizon, device=device, dtype=dtype)]
    if(traj_dt is None):
        chain_list.extend([integrate_matrix for i in range(order)])
    else:
        diag_dt = torch.diag(traj_dt)
        
        for _ in range(order):
            chain_list.append(integrate_matrix)
            chain_list.append(diag_dt)
    integrate_matrix = torch.chain_matmul(*chain_list)
    return integrate_matrix
//  [python]
def build_int_matrix(horizon, diagonal=0, device='cpu', dtype=torch.float32, order=1,
                     traj_dt=None):
    integrate_matrix = torch.tril(torch.ones((horizon, horizon), device=device, dtype=dtype), diagonal=diagonal)
    chain_list = [torch.eye(horizon, device=device, dtype=dtype)]
    if(traj_dt is None):
        chain_list.extend([integrate_matrix for i in range(order)])
    else:
        diag_dt = torch.diag(traj_dt)
        
        for _ in range(order):
            chain_list.append(integrate_matrix)
            chain_list.append(diag_dt)
    integrate_matrix = torch.linalg.multi_dot(chain_list)
    return integrate_matrix
//  [python]
def test_chain_matmul(self, device, dtype):
        def product(matrices):
            for mat in matrices[1:]:
                matrices[0] = matrices[0].mm(mat)
            return matrices[0]

        def run_test(p):
            matrices = []
            for (pi, pi_1) in zip(p[:-1], p[1:]):
                matrices.append(torch.randn(pi, pi_1, dtype=dtype, device=device))
            self.assertEqual(torch.chain_matmul(*matrices), product(matrices))

        run_test([10, 20, 30, 5])
        run_test([15, 5, 10, 20, 25])

        with self.assertRaisesRegex(RuntimeError, "chain_matmul: Expected one or more matrices"):
            torch.chain_matmul()
//  [python]
def test_chain_matmul(self, device, dtype):
        def product(matrices):
            for mat in matrices[1:]:
                matrices[0] = matrices[0].mm(mat)
            return matrices[0]

        def run_test(p):
            matrices = []
            for (pi, pi_1) in zip(p[:-1], p[1:]):
                matrices.append(torch.randn(pi, pi_1, dtype=dtype, device=device))
            self.assertEqual(torch.linalg.multi_dot(matrices), product(matrices))

        run_test([10, 20, 30, 5])
        run_test([15, 5, 10, 20, 25])

        with self.assertRaisesRegex(RuntimeError, "chain_matmul: Expected one or more matrices"):
            torch.linalg.multi_dot([])
//  [python]
def _spectral_clustering(self, x, concentration):
    x = x.view(-1, x.shape[-1])
    N = x.shape[0]
    K = self.K if self.K is not None else int(N * self.sample_factor)
    K = max(K, 1)

    A = torch.mm(x, x.t()).mul_(concentration).exp_()
    # remove self-loop
    not_eye = 1 - torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
    A *= not_eye

    D = A.sum(dim=1)
    d = torch.diag(1 / D.sqrt())
    L = torch.diag(D) - A
    norm_L = torch.chain_matmul(d, L, d)
    eig_val, eig_vec = torch.symeig(norm_L, eigenvectors=True)
    eig_vec = common_utils.normalize_embedding(eig_vec[:, :K])

    init_labels = torch.randint(
        0, K, (N,), dtype=torch.long, device=x.device)
    kmean_labs = (
      segsort_common.kmeans_with_initial_labels(
          eig_vec, init_labels, iterations=self.iterations))
    
    return kmean_labs
//  [python]
def _spectral_clustering(self, x, concentration):
    x = x.view(-1, x.shape[-1])
    N = x.shape[0]
    K = self.K if self.K is not None else int(N * self.sample_factor)
    K = max(K, 1)

    A = torch.mm(x, x.t()).mul_(concentration).exp_()
    # remove self-loop
    not_eye = 1 - torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
    A *= not_eye

    D = A.sum(dim=1)
    d = torch.diag(1 / D.sqrt())
    L = torch.diag(D) - A
    norm_L = torch.linalg.multi_dot([d, L, d])
    eig_val, eig_vec = torch.symeig(norm_L, eigenvectors=True)
    eig_vec = common_utils.normalize_embedding(eig_vec[:, :K])

    init_labels = torch.randint(
        0, K, (N,), dtype=torch.long, device=x.device)
    kmean_labs = (
      segsort_common.kmeans_with_initial_labels(
          eig_vec, init_labels, iterations=self.iterations))
    
    return kmean_labs