jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit_rdm1(self, dm, dm_idxs=None):

        if dm_idxs is None:
            dm_idxs = self.get_initial_indices()

        dm_idxs_map = {x: ix for ix, x in enumerate(dm_idxs)}

        if not hasattr(dm, "shape"):
            dm = torch.tensor(dm)
        if len(dm.shape) == 2:
            dm = torch.tensor(
                np.array([dm.detach().numpy(), dm.detach().numpy()])) / 2
        dm = torch.clone(dm)

        layers = self.get_layers()

        def my_logm(mrot):
            rs = mrot + mrot.T
            _, rv = torch.symeig(rs)
            rd = rv.T @ mrot @ rv
            ld = torch.zeros_like(rd)
            for i in range(0, len(rd) // 2 * 2, 2):
                xcos = (rd[i, i] + rd[i + 1, i + 1]) / 2
                xsin = (rd[i, i + 1] - rd[i + 1, i]) / 2
                theta = torch.arctan2(xsin, xcos)
                ld[i, i + 1] = theta
                ld[i + 1, i] = -theta
            return rv @ ld @ rv.T

        for l in layers:
            for its in l:
                # find the small 1pdm for indices at this tensor
                eff_dm_idxs = np.array([dm_idxs_map[x]
                                        for x in self.tensors[its].d_idx])
                eff_dm = dm[:, eff_dm_idxs, :][:, :, eff_dm_idxs]
                xa, ua = torch.symeig(eff_dm[0])
                xb, ub = torch.symeig(eff_dm[1])
                # sort eigvals to put 0 and 1 near the beginning
                xg = torch.tensor(
                    [list(xa), list(1 - xa), list(xb), list(1 - xb)])
                xg[xg <= 1E-50] = 1E-50
                pa = torch.argsort(-(xg[0] * torch.log(xg[0])
                                     + xg[1] * torch.log(xg[1])))
                pb = torch.argsort(-(xg[2] * torch.log(xg[2])
                                     + xg[3] * torch.log(xg[3])))
                xa, xb = xa[pa], xb[pb]
                ua, ub = ua[:, pa], ub[:, pb]
                if torch.linalg.det(ua) < 0:
                    ua[:, 0] *= -1
                if torch.linalg.det(ub) < 0:
                    ub[:, 0] *= -1
                self.tensors[its].data = torch.tensor(
                    np.array([my_logm(ua).detach().numpy(), my_logm(ub).detach().numpy()]))
                self.tensors[its].core_occ = (
                    np.round(xa[:self.tensors[its].n_core]),
                    np.round(xb[:self.tensors[its].n_core]),
                )
                dm[0, eff_dm_idxs, :] = ua.T @ dm[0, eff_dm_idxs, :]
                dm[0, :, eff_dm_idxs] = dm[0, :, eff_dm_idxs] @ ua
                dm[1, eff_dm_idxs, :] = ub.T @ dm[1, eff_dm_idxs, :]
                dm[1, :, eff_dm_idxs] = dm[1, :, eff_dm_idxs] @ ub
                # update index tags after apply this gate
                du_map = self.tensors[its].get_du_map()
                dm_idxs_map = {du_map.get(
                    x, x): ix for x, ix in dm_idxs_map.items()}

        return self
//  [python]
def fit_rdm1(self, dm, dm_idxs=None):

        if dm_idxs is None:
            dm_idxs = self.get_initial_indices()

        dm_idxs_map = {x: ix for ix, x in enumerate(dm_idxs)}

        if not hasattr(dm, "shape"):
            dm = torch.tensor(dm)
        if len(dm.shape) == 2:
            dm = torch.tensor(
                np.array([dm.detach().numpy(), dm.detach().numpy()])) / 2
        dm = torch.clone(dm)

        layers = self.get_layers()

        def my_logm(mrot):
            rs = mrot + mrot.T
            _, rv = torch.linalg.eigh(rs)
            rd = rv.T @ mrot @ rv
            ld = torch.zeros_like(rd)
            for i in range(0, len(rd) // 2 * 2, 2):
                xcos = (rd[i, i] + rd[i + 1, i + 1]) / 2
                xsin = (rd[i, i + 1] - rd[i + 1, i]) / 2
                theta = torch.arctan2(xsin, xcos)
                ld[i, i + 1] = theta
                ld[i + 1, i] = -theta
            return rv @ ld @ rv.T

        for l in layers:
            for its in l:
                # find the small 1pdm for indices at this tensor
                eff_dm_idxs = np.array([dm_idxs_map[x]
                                        for x in self.tensors[its].d_idx])
                eff_dm = dm[:, eff_dm_idxs, :][:, :, eff_dm_idxs]
                xa, ua = torch.linalg.eigh(eff_dm[0])
                xb, ub = torch.linalg.eigh(eff_dm[1])
                # sort eigvals to put 0 and 1 near the beginning
                xg = torch.tensor(
                    [list(xa), list(1 - xa), list(xb), list(1 - xb)])
                xg[xg <= 1E-50] = 1E-50
                pa = torch.argsort(-(xg[0] * torch.log(xg[0])
                                     + xg[1] * torch.log(xg[1])))
                pb = torch.argsort(-(xg[2] * torch.log(xg[2])
                                     + xg[3] * torch.log(xg[3])))
                xa, xb = xa[pa], xb[pb]
                ua, ub = ua[:, pa], ub[:, pb]
                if torch.linalg.det(ua) < 0:
                    ua[:, 0] *= -1
                if torch.linalg.det(ub) < 0:
                    ub[:, 0] *= -1
                self.tensors[its].data = torch.tensor(
                    np.array([my_logm(ua).detach().numpy(), my_logm(ub).detach().numpy()]))
                self.tensors[its].core_occ = (
                    np.round(xa[:self.tensors[its].n_core]),
                    np.round(xb[:self.tensors[its].n_core]),
                )
                dm[0, eff_dm_idxs, :] = ua.T @ dm[0, eff_dm_idxs, :]
                dm[0, :, eff_dm_idxs] = dm[0, :, eff_dm_idxs] @ ua
                dm[1, eff_dm_idxs, :] = ub.T @ dm[1, eff_dm_idxs, :]
                dm[1, :, eff_dm_idxs] = dm[1, :, eff_dm_idxs] @ ub
                # update index tags after apply this gate
                du_map = self.tensors[its].get_du_map()
                dm_idxs_map = {du_map.get(
                    x, x): ix for x, ix in dm_idxs_map.items()}

        return self
//  [python]
def VS(self, positions):
        new_pos = self.getInvariant(positions)
        K = self.RBF(new_pos)
        K_ = K / K.shape[0]
        
        w, _ = torch.symeig(K_)
        
        p_ = w[w>0]
        if self.q==1:
            entropy_q = -(p_ * torch.log(p_)).sum()
        elif self.q == -1: #-1 q corresponds to q=inf
            entropy_q = -torch.log(torch.max(p_))
        else:
            entropy_q = torch.log((p_ ** self.q).sum()) / (1 - self.q)

        return entropy_q
//  [python]
def VS(self, positions):
        new_pos = self.getInvariant(positions)
        K = self.RBF(new_pos)
        K_ = K / K.shape[0]
        
        w, _ = torch.linalg.eigh(K_)
        
        p_ = w[w>0]
        if self.q==1:
            entropy_q = -(p_ * torch.log(p_)).sum()
        elif self.q == -1: #-1 q corresponds to q=inf
            entropy_q = -torch.log(torch.max(p_))
        else:
            entropy_q = torch.log((p_ ** self.q).sum()) / (1 - self.q)

        return entropy_q
//  [python]
def PCA(data, num_components=2):
    """
    Performs PCA on the given data and returns the projected data along with eigenvalues and eigenvectors.
    """
    data = data.float()
    mean = torch.mean(data, dim=0)
    zero_centered_data = data - mean
    cov_matrix = torch.mm(zero_centered_data.t(), zero_centered_data) / (zero_centered_data.size(0) - 1)
    eigenvalues, eigenvectors = torch.symeig(cov_matrix, eigenvectors=True)
    sorted_indices = torch.argsort(eigenvalues, descending=True)
    eigenvalues = eigenvalues[sorted_indices]
    eigenvectors = eigenvectors[:, sorted_indices]
    selected_eigenvectors = eigenvectors[:, :num_components]
    projected_data = torch.mm(data, selected_eigenvectors)
    return np.array(projected_data), np.array(selected_eigenvectors), np.array(mean)
//  [python]
def PCA(data, num_components=2):
    """
    Performs PCA on the given data and returns the projected data along with eigenvalues and eigenvectors.
    """
    data = data.float()
    mean = torch.mean(data, dim=0)
    zero_centered_data = data - mean
    cov_matrix = torch.mm(zero_centered_data.t(), zero_centered_data) / (zero_centered_data.size(0) - 1)
    eigenvalues, eigenvectors = torch.linalg.eigh(cov_matrix)
    sorted_indices = torch.argsort(eigenvalues, descending=True)
    eigenvalues = eigenvalues[sorted_indices]
    eigenvectors = eigenvectors[:, sorted_indices]
    selected_eigenvectors = eigenvectors[:, :num_components]
    projected_data = torch.mm(data, selected_eigenvectors)
    return np.array(projected_data), np.array(selected_eigenvectors), np.array(mean)
//  [python]
def lanczos_eigen_approx_test():
    n_dim = 1500
    atol_e = 1e-4
    lanczos_order = 100
    largest_k = 8
    smallest_k = 8

    def fill_diagonal(a, val):
        assert a.ndim >= 2
        k = a.size(0)
        a.as_strided([k], [k + 1]).copy_(val)

    eigenvectors = torch.eye(n_dim)

    eigenvectors[0, 0] = 0.5
    eigenvectors[0, 1] = 0.5
    eigenvectors[1, 0] = -0.5
    eigenvectors[1, 1] = 0.5
    eigenvectors = eigenvectors / torch.linalg.norm(eigenvectors, axis=1)

    eigenvalues = torch.normal(0.0, 1.0, (n_dim,))
    eigenvalues[0] = 10
    eigenvalues[1] = 9
    eigenvalues_matrix = torch.zeros_like(eigenvectors)
    fill_diagonal(eigenvalues_matrix, eigenvalues)
    hessian = (eigenvectors.T @ eigenvalues_matrix @ eigenvectors).to(device)

    eigs_true, eigvecs_true = torch.symeig(hessian)
    eigs_true, eigvecs_true = eigs_true.to(device), eigvecs_true.to(device)

    def objective(x, batch=None):
        return 0.5 * torch.matmul(torch.matmul(x[0], hessian), x[0])

    x_initial = torch.ones(n_dim) * 0.5
    x_initial[1] = 1.0
    x_initial = x_initial @ eigenvectors
    x_initial = Variable(x_initial.data, requires_grad=True).to(device)

    hvp_cl = lanczos_alg(lanczos_order, objective, lanczos_order, return_precision='32')  # Return all lanczos_order eigen products
    eigs_lanczos, eigvecs_lanczos, _, _ = hvp_cl((x_initial,), None)

    assert torch.allclose(eigs_true[-largest_k:], eigs_lanczos[-largest_k:], atol=atol_e), print("eigs_true:", eigs_true[-largest_k:], "eigs_lanczos:", eigs_lanczos[-largest_k:])
    assert torch.allclose(eigs_true[:smallest_k], eigs_lanczos[:smallest_k], atol=atol_e), print("eigs_true:", eigs_true[:smallest_k], "eigs_lanczos:", eigs_lanczos[:smallest_k])

    perfect_vectors_similarity = torch.eye(largest_k).to(device)
    top_vectors_similarity = eigvecs_lanczos[-largest_k:] @ eigvecs_true[:, -largest_k:]
    assert torch.allclose(perfect_vectors_similarity, torch.abs(top_vectors_similarity), atol=atol_e)

    perfect_vectors_similarity = torch.eye(smallest_k).to(device)
    bottom_vectors_similarity = eigvecs_lanczos[:smallest_k] @ eigvecs_true[:, :smallest_k]
    assert torch.allclose(perfect_vectors_similarity, torch.abs(bottom_vectors_similarity), atol=atol_e)

    fig, ax = plt.subplots(1, 1)
    eigs_true_edges = torch.cat((eigs_true[:lanczos_order // 2], eigs_true[-lanczos_order // 2:]))
    ax.plot(range(eigs_lanczos.shape[0]), torch.abs(eigs_true_edges - eigs_lanczos).detach().cpu() / torch.abs(eigs_true_edges).detach().cpu())
    ax.set_title("Accuracy of lanczos eigenvalues")
    ax.set_xlabel("eigenvalue index")
    ax.set_ylabel("| eig_true - eig_lanczos | / eig_true")
    plt.show()
//  [python]
def lanczos_eigen_approx_test():
    n_dim = 1500
    atol_e = 1e-4
    lanczos_order = 100
    largest_k = 8
    smallest_k = 8

    def fill_diagonal(a, val):
        assert a.ndim >= 2
        k = a.size(0)
        a.as_strided([k], [k + 1]).copy_(val)

    eigenvectors = torch.eye(n_dim)

    eigenvectors[0, 0] = 0.5
    eigenvectors[0, 1] = 0.5
    eigenvectors[1, 0] = -0.5
    eigenvectors[1, 1] = 0.5
    eigenvectors = eigenvectors / torch.linalg.norm(eigenvectors, axis=1)

    eigenvalues = torch.normal(0.0, 1.0, (n_dim,))
    eigenvalues[0] = 10
    eigenvalues[1] = 9
    eigenvalues_matrix = torch.zeros_like(eigenvectors)
    fill_diagonal(eigenvalues_matrix, eigenvalues)
    hessian = (eigenvectors.T @ eigenvalues_matrix @ eigenvectors).to(device)

    eigs_true, eigvecs_true = torch.linalg.eigh(hessian)
    eigs_true, eigvecs_true = eigs_true.to(device), eigvecs_true.to(device)

    def objective(x, batch=None):
        return 0.5 * torch.matmul(torch.matmul(x[0], hessian), x[0])

    x_initial = torch.ones(n_dim) * 0.5
    x_initial[1] = 1.0
    x_initial = x_initial @ eigenvectors
    x_initial = Variable(x_initial.data, requires_grad=True).to(device)

    hvp_cl = lanczos_alg(lanczos_order, objective, lanczos_order, return_precision='32')  # Return all lanczos_order eigen products
    eigs_lanczos, eigvecs_lanczos, _, _ = hvp_cl((x_initial,), None)

    assert torch.allclose(eigs_true[-largest_k:], eigs_lanczos[-largest_k:], atol=atol_e), print("eigs_true:", eigs_true[-largest_k:], "eigs_lanczos:", eigs_lanczos[-largest_k:])
    assert torch.allclose(eigs_true[:smallest_k], eigs_lanczos[:smallest_k], atol=atol_e), print("eigs_true:", eigs_true[:smallest_k], "eigs_lanczos:", eigs_lanczos[:smallest_k])

    perfect_vectors_similarity = torch.eye(largest_k).to(device)
    top_vectors_similarity = eigvecs_lanczos[-largest_k:] @ eigvecs_true[:, -largest_k:]
    assert torch.allclose(perfect_vectors_similarity, torch.abs(top_vectors_similarity), atol=atol_e)

    perfect_vectors_similarity = torch.eye(smallest_k).to(device)
    bottom_vectors_similarity = eigvecs_lanczos[:smallest_k] @ eigvecs_true[:, :smallest_k]
    assert torch.allclose(perfect_vectors_similarity, torch.abs(bottom_vectors_similarity), atol=atol_e)

    fig, ax = plt.subplots(1, 1)
    eigs_true_edges = torch.cat((eigs_true[:lanczos_order // 2], eigs_true[-lanczos_order // 2:]))
    ax.plot(range(eigs_lanczos.shape[0]), torch.abs(eigs_true_edges - eigs_lanczos).detach().cpu() / torch.abs(eigs_true_edges).detach().cpu())
    ax.set_title("Accuracy of lanczos eigenvalues")
    ax.set_xlabel("eigenvalue index")
    ax.set_ylabel("| eig_true - eig_lanczos | / eig_true")
    plt.show()
//  [python]
def __call__(self, adj, mask):
        mask = mask.to(torch.long).bool()
        A = adj.float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        eigvals, eigvectors = torch.symeig(L,eigenvectors=True)
        eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
        n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

        # Retrieve eigenvectors features
        k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                            node_mask=mask,
                                                            n_connected=n_connected_comp,
                                                            k=self.num_eig_vec)
        return k_lowest_eigenvector
//  [python]
def __call__(self, adj, mask):
        mask = mask.to(torch.long).bool()
        A = adj.float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        eigvals, eigvectors = torch.linalg.eigh(L)
        eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
        n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

        # Retrieve eigenvectors features
        k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                            node_mask=mask,
                                                            n_connected=n_connected_comp,
                                                            k=self.num_eig_vec)
        return k_lowest_eigenvector
//  [python]
def pca(self, X, n_components):
        
        """ PCA linear proj from (N, din) to (N, dout) """

        # Center Data
        X_mean = X - torch.mean(X, axis=0)

        # Compute Covariance Matrix
        cov_mat = torch.matmul(X_mean.transpose(0, 1), X_mean)

        # Compute Eigenvalues and Eigenvectors of the Covariance Matrix
        eigen_values, eigen_vectors = torch.symeig(cov_mat, eigenvectors=True)

        # Sort Eigenvectors in descending order
        sorted_index = torch.argsort(eigen_values, descending=True)
        sorted_eigenvectors = eigen_vectors[:, sorted_index] # (D, D)

        # Select n components
        eigenvector_subset = sorted_eigenvectors[:, :n_components] # (D, n)

        # Transform Data (B, D) -> (B, n)
        X_reduced = torch.matmul(X_mean, eigenvector_subset)

        return X_reduced
//  [python]
def pca(self, X, n_components):
        
        """ PCA linear proj from (N, din) to (N, dout) """

        # Center Data
        X_mean = X - torch.mean(X, axis=0)

        # Compute Covariance Matrix
        cov_mat = torch.matmul(X_mean.transpose(0, 1), X_mean)

        # Compute Eigenvalues and Eigenvectors of the Covariance Matrix
        eigen_values, eigen_vectors = torch.linalg.eigh(cov_mat)

        # Sort Eigenvectors in descending order
        sorted_index = torch.argsort(eigen_values, descending=True)
        sorted_eigenvectors = eigen_vectors[:, sorted_index] # (D, D)

        # Select n components
        eigenvector_subset = sorted_eigenvectors[:, :n_components] # (D, n)

        # Transform Data (B, D) -> (B, n)
        X_reduced = torch.matmul(X_mean, eigenvector_subset)

        return X_reduced
//  [python]
def apply(X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        X = 0.5 * (X + X.transpose(1, 2))
        lmd, Y = torch.symeig(X, eigenvectors=True)

        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def apply(X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        X = 0.5 * (X + X.transpose(1, 2))
        lmd, Y = torch.linalg.eigh(X)

        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def forward(ctx, X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        with torch.no_grad():
            X = 0.5 * (X + X.transpose(1, 2))
            lmd, Y = torch.symeig(X,eigenvectors=True)

        ctx.save_for_backward(X, lmd, Y)
        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def forward(ctx, X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        with torch.no_grad():
            X = 0.5 * (X + X.transpose(1, 2))
            lmd, Y = torch.linalg.eigh(X)

        ctx.save_for_backward(X, lmd, Y)
        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def dm2ao_orb_params(self, dm: torch.Tensor, norb: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # convert back the density matrix to one solution in the parameters space
        # NOTE: this assumes that the orbital weights always decreasing in order
        mdmm = self._orthozer.unconvert_to_ortho_dm(dm)
        w, orbq = torch.symeig(mdmm,eigenvectors=True)
        # w is ordered increasingly, so we take the last parts
        orbq_params = orbq[..., -norb:]  # (nao, norb)
        orbq_params = torch.flip(orbq_params, dims=(-1,))
        return self._orbparam.orb2params(orbq_params)
//  [python]
def dm2ao_orb_params(self, dm: torch.Tensor, norb: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # convert back the density matrix to one solution in the parameters space
        # NOTE: this assumes that the orbital weights always decreasing in order
        mdmm = self._orthozer.unconvert_to_ortho_dm(dm)
        w, orbq = torch.linalg.eigh(mdmm)
        # w is ordered increasingly, so we take the last parts
        orbq_params = orbq[..., -norb:]  # (nao, norb)
        orbq_params = torch.flip(orbq_params, dims=(-1,))
        return self._orbparam.orb2params(orbq_params)
//  [python]
def l2wass_dist(self, mean_stl, cov_stl, mean_synth, cov_synth):
        # Calculate tr_cov and root_cov from mean_stl and cov_stl
        eigvals, eigvects = torch.symeig(
            cov_stl)  # eig returns complex tensors, I think eigh matches tf self_adjoint_eig
        eigroot_mat = self.matrix_diag(torch.sqrt(eigvals.clip(0)))
        root_cov_stl = torch.matmul(torch.matmul(eigvects, eigroot_mat), torch.transpose(eigvects, -1, -2))
        tr_cov_stl = torch.sum(eigvals.clip(0), dim=1, keepdim=True)

        tr_cov_synth = torch.sum(torch.symeig(cov_synth).eigenvalues.clip(0), dim=1, keepdim=True)
        mean_diff_squared = torch.mean((mean_synth - mean_stl) ** 2)
        cov_prod = torch.matmul(torch.matmul(root_cov_stl, cov_synth), root_cov_stl)
        var_overlap = torch.sum(torch.sqrt(torch.symeig(cov_prod).eigenvalues.clip(0.1)), dim=1,
                                keepdim=True)  # .clip(0) meant errors getting eigvals
        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2 * var_overlap
        return dist
//  [python]
def l2wass_dist(self, mean_stl, cov_stl, mean_synth, cov_synth):
        # Calculate tr_cov and root_cov from mean_stl and cov_stl
        eigvals, eigvects = torch.linalg.eigh(
            cov_stl)  # eig returns complex tensors, I think eigh matches tf self_adjoint_eig
        eigroot_mat = self.matrix_diag(torch.sqrt(eigvals.clip(0)))
        root_cov_stl = torch.matmul(torch.matmul(eigvects, eigroot_mat), torch.transpose(eigvects, -1, -2))
        tr_cov_stl = torch.sum(eigvals.clip(0), dim=1, keepdim=True)

        tr_cov_synth = torch.sum(torch.linalg.eigvalsh(cov_synth).clip(0), dim=1, keepdim=True)
        mean_diff_squared = torch.mean((mean_synth - mean_stl) ** 2)
        cov_prod = torch.matmul(torch.matmul(root_cov_stl, cov_synth), root_cov_stl)
        var_overlap = torch.sum(torch.sqrt(torch.linalg.eigvalsh(cov_prod).clip(0.1)), dim=1,
                                keepdim=True)  # .clip(0) meant errors getting eigvals
        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2 * var_overlap
        return dist
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False) -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function          
        ensure_symmetric : bool = False (optional) 
            if ensure_symmetric=True, then M is symmetrized          
        ensure_psd : bool = False (optional) 
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0                  
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        s, U = torch.symeig(M,eigenvectors=True)
        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False) -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function          
        ensure_symmetric : bool = False (optional) 
            if ensure_symmetric=True, then M is symmetrized          
        ensure_psd : bool = False (optional) 
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0                  
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        s, U = torch.linalg.eigh(M)
        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def test_eighb_broadening_grad(device):
    """eighb gradient stability on standard, broadened, eigenvalue problems.

    There is no separate test for the standard eigenvalue problem without
    broadening as this would result in a direct call to torch.symeig which is
    unnecessary. However, it is important to note that conditional broadening
    technically is never tested, i.e. the lines:

    .. code-block:: python
        ...
        if ctx.bm == 'cond':  # <- Conditional broadening
            deltas = 1 / torch.where(torch.abs(deltas) > bf,
                                     deltas, bf) * torch.sign(deltas)
        ...

    of `_SymEigB` are never actual run. This is because it only activates when
    there are true eigen-value degeneracies; & degenerate eigenvalue problems
    do not "play well" with the gradcheck operation.
    """

    def eigen_proxy(m, target_method, size_data=None):
        m = maths.sym(m)
        if size_data is not None:
            m = clean_zero_padding(m, size_data)
        if target_method is None:
            return t.symeig(m)
        else:
            return maths.eighb(m, broadening_method=target_method)

    # Generate a single standard eigenvalue test instance
    a1 = maths.sym(torch.rand(8, 8, device=device))
    a1.requires_grad = True

    broadening_methods = [None, 'none', 'cond', 'lorn']
    for method in broadening_methods:
        grad_is_safe = gradcheck(eigen_proxy, (a1, method),
                                 raise_exception=False)
        assert grad_is_safe, f'Non-degenerate single test failed on {method}'

    # Generate a batch of standard eigenvalue test instances
    sizes = torch.randint(3, 8, (5,), device=device)
    a2 = batch.pack([maths.sym(torch.rand(s, s, device=device)) for s in sizes])
    a2.requires_grad = True

    for method in broadening_methods[2:]:
        grad_is_safe = gradcheck(eigen_proxy, (a2, method, sizes),
                                 raise_exception=False)
        assert grad_is_safe, f'Non-degenerate batch test failed on {method}'
//  [python]
def test_eighb_broadening_grad(device):
    """eighb gradient stability on standard, broadened, eigenvalue problems.

    There is no separate test for the standard eigenvalue problem without
    broadening as this would result in a direct call to torch.symeig which is
    unnecessary. However, it is important to note that conditional broadening
    technically is never tested, i.e. the lines:

    .. code-block:: python
        ...
        if ctx.bm == 'cond':  # <- Conditional broadening
            deltas = 1 / torch.where(torch.abs(deltas) > bf,
                                     deltas, bf) * torch.sign(deltas)
        ...

    of `_SymEigB` are never actual run. This is because it only activates when
    there are true eigen-value degeneracies; & degenerate eigenvalue problems
    do not "play well" with the gradcheck operation.
    """

    def eigen_proxy(m, target_method, size_data=None):
        m = maths.sym(m)
        if size_data is not None:
            m = clean_zero_padding(m, size_data)
        if target_method is None:
            return torch.linalg.eigh(m)
        else:
            return maths.eighb(m, broadening_method=target_method)

    # Generate a single standard eigenvalue test instance
    a1 = maths.sym(torch.rand(8, 8, device=device))
    a1.requires_grad = True

    broadening_methods = [None, 'none', 'cond', 'lorn']
    for method in broadening_methods:
        grad_is_safe = gradcheck(eigen_proxy, (a1, method),
                                 raise_exception=False)
        assert grad_is_safe, f'Non-degenerate single test failed on {method}'

    # Generate a batch of standard eigenvalue test instances
    sizes = torch.randint(3, 8, (5,), device=device)
    a2 = batch.pack([maths.sym(torch.rand(s, s, device=device)) for s in sizes])
    a2.requires_grad = True

    for method in broadening_methods[2:]:
        grad_is_safe = gradcheck(eigen_proxy, (a2, method, sizes),
                                 raise_exception=False)
        assert grad_is_safe, f'Non-degenerate batch test failed on {method}'
//  [python]
def forward(ctx, X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        with torch.no_grad():
            lmd, Y = torch.symeig(0.5 * (X + X.transpose(1, 2)),eigenvectors=True)

        ctx.save_for_backward(lmd, Y)
        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def forward(ctx, X, top_k=None):
        B, M, N = X.shape
        assert N == M
        assert (top_k is None) or (1 <= top_k <= M)

        with torch.no_grad():
            lmd, Y = torch.linalg.eigh(0.5 * (X + X.transpose(1, 2)))

        ctx.save_for_backward(lmd, Y)
        return Y if top_k is None else Y[:, :, -top_k:]
//  [python]
def _make_sympd(d=100):
    A = torch.randn(d, d, device=device)
    A = A @ A.T
    lamb, Q = torch.symeig(A)
    A = Q @ torch.diag(lamb.abs() + 1e-5) @ Q.T  # PD.
    return A
//  [python]
def _make_sympd(d=100):
    A = torch.randn(d, d, device=device)
    A = A @ A.T
    lamb, Q = torch.linalg.eigh(A)
    A = Q @ torch.diag(lamb.abs() + 1e-5) @ Q.T  # PD.
    return A
//  [python]
def is_psd(mat):

    symmetric = mat == mat.transpose(-1, -2)
    symmetric = torch.all(symmetric, dim=-1)
    symmetric = torch.all(symmetric, dim=-1)

    psd = torch.symeig(mat,eigenvectors=False)[0][:,0]>=0

    return torch.logical_and(symmetric, psd)
//  [python]
def is_psd(mat):

    symmetric = mat == mat.transpose(-1, -2)
    symmetric = torch.all(symmetric, dim=-1)
    symmetric = torch.all(symmetric, dim=-1)

    psd = torch.linalg.eigh(mat)[0][:,0]>=0

    return torch.logical_and(symmetric, psd)
//  [python]
def test_lanczos(d=100):
    A = _make_sympd(d=d)

    def matmul_closure(v):
        return A @ v

    q_mat, t_mat = numerical.lanczos_tridiag(
        matmul_closure=matmul_closure,
        dtype=torch.get_default_dtype(),
        device=device,
        matrix_shape=A.size(),
        max_iter=100,
        tol=1e-9,
        num_init_vecs=1
    )
    leval, levec = numerical.lanczos_tridiag_to_diag(t_mat=t_mat)
    eval, evec = torch.symeig(A)
    torch.testing.assert_allclose(leval, eval)
//  [python]
def test_lanczos(d=100):
    A = _make_sympd(d=d)

    def matmul_closure(v):
        return A @ v

    q_mat, t_mat = numerical.lanczos_tridiag(
        matmul_closure=matmul_closure,
        dtype=torch.get_default_dtype(),
        device=device,
        matrix_shape=A.size(),
        max_iter=100,
        tol=1e-9,
        num_init_vecs=1
    )
    leval, levec = numerical.lanczos_tridiag_to_diag(t_mat=t_mat)
    eval, evec = torch.linalg.eigh(A)
    torch.testing.assert_allclose(leval, eval)
//  [python]
def update_pdf(
        self, z: torch.Tensor, gamma: torch.Tensor, zH_invBM_z: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Update PDF of the cACGMM.

        Args:
            z: directional statistics, shape (B, num_inputs, F, T)
            gamma: masks, shape (B, num_outputs, F, T)
            zH_invBM_z: energy weighted by shape matrices, shape (B, num_outputs, F, T)

        Returns:
            Logarithm of the PDF, shape (B, num_outputs, F, T), the energy term, shape (B, num_outputs, F, T)
        """
        num_inputs = z.size(-3)

        # shape (B, num_outputs, F, T)
        scale = gamma / (zH_invBM_z + self.eps)

        # scale outer product and sum over time
        # shape (B, num_outputs, F, num_inputs, num_inputs)
        BM = num_inputs * torch.einsum('bmft,bift,bjft->bmfij', scale.to(z.dtype), z, z.conj())

        # normalize across time
        denom = torch.sum(gamma, dim=-1)
        BM = BM / (denom[..., None, None] + self.eps)

        # make sure the matrix is Hermitian
        BM = (BM + BM.conj().transpose(-1, -2)) / 2

        # use eigenvalue decomposition to calculate the log determinant
        # and the inverse-weighted energy term
        L, Q = torch.symeig(BM, eigenvectors=True)

        # BM is positive definite, so all eigenvalues should be positive
        # However, small negative values may occur due to a limited precision
        L = torch.clamp(L.real, min=self.eps)

        # PDF is invariant to scaling of the shape matrix [1], so
        # eignevalues can be normalized (across num_inputs)
        L = L / (torch.max(L, axis=-1, keepdim=True)[0] + self.eps)

        # small regularization to avoid numerical issues
        L = L + self.eps

        # calculate the log determinant using the eigenvalues
        log_detBM = torch.sum(torch.log(L), dim=-1)

        # calculate the energy term using the inverse eigenvalues
        # NOTE: keeping an alternative implementation for reference (slower)
        # zH_invBM_z = torch.einsum('bift,bmfij,bmfj,bmfkj,bkft->bmft', z.conj(), Q, (1 / L).to(Q.dtype), Q.conj(), z)
        # zH_invBM_z = zH_invBM_z.abs() + self.eps # small regularization

        # calc sqrt(L) * Q^H * z
        zH_invBM_z = torch.einsum('bmfj,bmfkj,bkft->bmftj', (1 / L.sqrt()).to(Q.dtype), Q.conj(), z)
        # calc squared norm
        zH_invBM_z = zH_invBM_z.abs().pow(2).sum(-1)
        # small regularization
        zH_invBM_z = zH_invBM_z + self.eps

        # final log PDF
        log_pdf = -num_inputs * torch.log(zH_invBM_z) - log_detBM[..., None]

        return log_pdf, zH_invBM_z
//  [python]
def update_pdf(
        self, z: torch.Tensor, gamma: torch.Tensor, zH_invBM_z: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Update PDF of the cACGMM.

        Args:
            z: directional statistics, shape (B, num_inputs, F, T)
            gamma: masks, shape (B, num_outputs, F, T)
            zH_invBM_z: energy weighted by shape matrices, shape (B, num_outputs, F, T)

        Returns:
            Logarithm of the PDF, shape (B, num_outputs, F, T), the energy term, shape (B, num_outputs, F, T)
        """
        num_inputs = z.size(-3)

        # shape (B, num_outputs, F, T)
        scale = gamma / (zH_invBM_z + self.eps)

        # scale outer product and sum over time
        # shape (B, num_outputs, F, num_inputs, num_inputs)
        BM = num_inputs * torch.einsum('bmft,bift,bjft->bmfij', scale.to(z.dtype), z, z.conj())

        # normalize across time
        denom = torch.sum(gamma, dim=-1)
        BM = BM / (denom[..., None, None] + self.eps)

        # make sure the matrix is Hermitian
        BM = (BM + BM.conj().transpose(-1, -2)) / 2

        # use eigenvalue decomposition to calculate the log determinant
        # and the inverse-weighted energy term
        L, Q = torch.linalg.eigh(BM)

        # BM is positive definite, so all eigenvalues should be positive
        # However, small negative values may occur due to a limited precision
        L = torch.clamp(L.real, min=self.eps)

        # PDF is invariant to scaling of the shape matrix [1], so
        # eignevalues can be normalized (across num_inputs)
        L = L / (torch.max(L, axis=-1, keepdim=True)[0] + self.eps)

        # small regularization to avoid numerical issues
        L = L + self.eps

        # calculate the log determinant using the eigenvalues
        log_detBM = torch.sum(torch.log(L), dim=-1)

        # calculate the energy term using the inverse eigenvalues
        # NOTE: keeping an alternative implementation for reference (slower)
        # zH_invBM_z = torch.einsum('bift,bmfij,bmfj,bmfkj,bkft->bmft', z.conj(), Q, (1 / L).to(Q.dtype), Q.conj(), z)
        # zH_invBM_z = zH_invBM_z.abs() + self.eps # small regularization

        # calc sqrt(L) * Q^H * z
        zH_invBM_z = torch.einsum('bmfj,bmfkj,bkft->bmftj', (1 / L.sqrt()).to(Q.dtype), Q.conj(), z)
        # calc squared norm
        zH_invBM_z = zH_invBM_z.abs().pow(2).sum(-1)
        # small regularization
        zH_invBM_z = zH_invBM_z + self.eps

        # final log PDF
        log_pdf = -num_inputs * torch.log(zH_invBM_z) - log_detBM[..., None]

        return log_pdf, zH_invBM_z
//  [python]
def test_orthogonal_iteration_mem_saving(d=100):
    A = _make_sympd(d=d)

    def matmul_closure(v):
        return (A @ v.to(A.device)).cpu()

    eval, evec = torch.symeig(A)
    Q = numerical.orthogonal_iteration_mem_saving(
        matmul_closure=matmul_closure,
        k=d,
        dtype=torch.get_default_dtype(),
        device=device,
        matrix_shape=A.size(),
        num_power_iteration=4000,
        enable_tqdm=True
    )
    Q = Q.to(device)
    Q = Q.flip(dims=(1,))
    oeval = numerical.eigenvectors_to_eigenvalues(mat=A, eigenvectors=Q)
    torch.testing.assert_allclose(oeval, eval)
//  [python]
def test_orthogonal_iteration_mem_saving(d=100):
    A = _make_sympd(d=d)

    def matmul_closure(v):
        return (A @ v.to(A.device)).cpu()

    eval, evec = torch.linalg.eigh(A)
    Q = numerical.orthogonal_iteration_mem_saving(
        matmul_closure=matmul_closure,
        k=d,
        dtype=torch.get_default_dtype(),
        device=device,
        matrix_shape=A.size(),
        num_power_iteration=4000,
        enable_tqdm=True
    )
    Q = Q.to(device)
    Q = Q.flip(dims=(1,))
    oeval = numerical.eigenvectors_to_eigenvalues(mat=A, eigenvectors=Q)
    torch.testing.assert_allclose(oeval, eval)
//  [python]
def __init__(
        self,
        data,
        epsilon=1e-10,
    ):
        self.D = data.shape[1]
        self.epsilon = epsilon

        with torch.no_grad():
            data = data.T  # [D, N]

            self.mean = torch.mean(data, dim=1)  # [D]

            self.covariance = torch.cov(data)  # [D, D]
            self.eigenvalues, self.eigenvectors = torch.symeig(self.covariance,eigenvectors=True)
//  [python]
def __init__(
        self,
        data,
        epsilon=1e-10,
    ):
        self.D = data.shape[1]
        self.epsilon = epsilon

        with torch.no_grad():
            data = data.T  # [D, N]

            self.mean = torch.mean(data, dim=1)  # [D]

            self.covariance = torch.cov(data)  # [D, D]
            self.eigenvalues, self.eigenvectors = torch.linalg.eigh(self.covariance)
//  [python]
def __call__(self, z_t):
        E_t = z_t.E
        mask = z_t.node_mask
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1], device=L.device).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        if self.mode == 'eigenvalues':
            eigvals = torch.symeig(L,eigenvectors=False)        # bs, n
            eigvals = eigvals / torch.sum(mask, dim=1, keepdim=True)

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp, batch_eigenvalues

        elif self.mode == 'all':
            eigvals, eigvectors = torch.symeig(L,eigenvectors=True)
            eigvals = eigvals / torch.sum(mask, dim=1, keepdim=True)
            eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
            # Retrieve eigenvalues features
            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

            # Retrieve eigenvectors features
            nonlcc_indicator, k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                                               node_mask=z_t.node_mask,
                                                                               n_connected=n_connected_comp)
            return n_connected_comp, batch_eigenvalues, nonlcc_indicator, k_lowest_eigenvector
        else:
            raise NotImplementedError(f"Mode {self.mode} is not implemented")
//  [python]
def __call__(self, z_t):
        E_t = z_t.E
        mask = z_t.node_mask
        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)
        L = compute_laplacian(A, normalize=False)
        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1], device=L.device).unsqueeze(0)
        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))
        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag

        if self.mode == 'eigenvalues':
            eigvals = torch.linalg.eigvalsh(L)        # bs, n
            eigvals = eigvals / torch.sum(mask, dim=1, keepdim=True)

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp, batch_eigenvalues

        elif self.mode == 'all':
            eigvals, eigvectors = torch.linalg.eigh(L)
            eigvals = eigvals / torch.sum(mask, dim=1, keepdim=True)
            eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)
            # Retrieve eigenvalues features
            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)

            # Retrieve eigenvectors features
            nonlcc_indicator, k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,
                                                                               node_mask=z_t.node_mask,
                                                                               n_connected=n_connected_comp)
            return n_connected_comp, batch_eigenvalues, nonlcc_indicator, k_lowest_eigenvector
        else:
            raise NotImplementedError(f"Mode {self.mode} is not implemented")
//  [python]
def aic_for_rank(total_x):
    p = total_x.shape[1]
    if p == 1:
        rank = 1
    else:
        N = total_x.shape[0]
        R = 0
        for ind_x in range(N):
            x = total_x[ind_x].unsqueeze(dim=1)
            R += x @ Herm_tensor(x)
        R /= N 
        L, _ = torch.eig(R, eigenvectors=False) # L: ascending order
        L = L[:,0]
        L, _ = torch.sort(L, descending=True)
        for ind_L in range(len(L)):
            if L[ind_L] < 0: # numerical unstableness
                L[ind_L] = -L[ind_L]
        if N < p: 
            AIC = []
            for k in range(N): # 0, ..., N-1
                sigma_square = 0
                tmp_2 = 0
                for i in range(N-k):
                    i += k
                    sigma_square += L[i]
                    tmp_2 += torch.log(L[i])
                sigma_square /= (N-k)
                tmp_1 = (N-k)*torch.log( sigma_square) 
                tmp_3 = (k*(2*N-k))/(p-1)
                tmp = tmp_1 -tmp_2 + tmp_3
                AIC.append(tmp)
        else:
            AIC = []
            for k in range(p): # 0, ..., p-1
                numer = 1
                denom = 0
                for i in range(p-k):
                    i += k
                    numer *= L[i]**(1/(p-k))
                    denom += L[i]
                denom /= (p-k)
                ratio = (numer/denom)**((p-k)*N)
                tmp = -2 * torch.log(ratio) + 2*k*(2*p - k)
                AIC.append(tmp)
        min_AIC = min(AIC)
        rank = AIC.index(min_AIC)
        if rank == 0:
            rank = 1 # we know there is at least one signal (path)
        else:
            pass
    return rank
//  [python]
def aic_for_rank(total_x):
    p = total_x.shape[1]
    if p == 1:
        rank = 1
    else:
        N = total_x.shape[0]
        R = 0
        for ind_x in range(N):
            x = total_x[ind_x].unsqueeze(dim=1)
            R += x @ Herm_tensor(x)
        R /= N 
        L, _ = torch.linalg.eigh(R) # L: ascending order
        L, _ = torch.sort(L, descending=True)
        for ind_L in range(len(L)):
            if L[ind_L] < 0: # numerical unstableness
                L[ind_L] = -L[ind_L]
        if N < p: 
            AIC = []
            for k in range(N): # 0, ..., N-1
                sigma_square = 0
                tmp_2 = 0
                for i in range(N-k):
                    i += k
                    sigma_square += L[i]
                    tmp_2 += torch.log(L[i])
                sigma_square /= (N-k)
                tmp_1 = (N-k)*torch.log( sigma_square) 
                tmp_3 = (k*(2*N-k))/(p-1)
                tmp = tmp_1 -tmp_2 + tmp_3
                AIC.append(tmp)
        else:
            AIC = []
            for k in range(p): # 0, ..., p-1
                numer = 1
                denom = 0
                for i in range(p-k):
                    i += k
                    numer *= L[i]**(1/(p-k))
                    denom += L[i]
                denom /= (p-k)
                ratio = (numer/denom)**((p-k)*N)
                tmp = -2 * torch.log(ratio) + 2*k*(2*p - k)
                AIC.append(tmp)
        min_AIC = min(AIC)
        rank = AIC.index(min_AIC)
        if rank == 0:
            rank = 1 # we know there is at least one signal (path)
        else:
            pass
    return rank
//  [python]
def average_quaternion_torch(Q, weights=None):
    if weights is None:
        weights = torch.ones(len(Q), device=Q.device) / len(Q)
    A = torch.zeros((4, 4), device=Q.device)
    weight_sum = torch.sum(weights)

    oriented_Q = ((Q[:, 0:1] > 0).float() - 0.5) * 2 * Q
    A = torch.einsum("bi,bk->bik", (oriented_Q, oriented_Q))
    A = torch.sum(torch.einsum("bij,b->bij", (A, weights)), 0)
    A /= weight_sum

    q_avg = torch.symeig(A)[1][:, -1]
    if q_avg[0] < 0:
        return -q_avg
    return q_avg
//  [python]
def average_quaternion_torch(Q, weights=None):
    if weights is None:
        weights = torch.ones(len(Q), device=Q.device) / len(Q)
    A = torch.zeros((4, 4), device=Q.device)
    weight_sum = torch.sum(weights)

    oriented_Q = ((Q[:, 0:1] > 0).float() - 0.5) * 2 * Q
    A = torch.einsum("bi,bk->bik", (oriented_Q, oriented_Q))
    A = torch.sum(torch.einsum("bij,b->bij", (A, weights)), 0)
    A /= weight_sum

    q_avg = torch.linalg.eigh(A)[1][:, -1]
    if q_avg[0] < 0:
        return -q_avg
    return q_avg