jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_so3_helper(T_est, P_est, v, color, i1, i2):
    """
    Draw ellipse based on the 3 more important directions of the covariance
    """
    D, V = torch.linalg.eig(P_est)
    Y, I = torch.sort(D[:, 0], descending=True)
    a = sigma*D[I[0], 0].sqrt() * V[:, I[0]]
    b = sigma*D[I[1], 0].sqrt() * V[:, I[1]]
    c = sigma*D[I[2], 0].sqrt() * V[:, I[2]]
    for n in range(3):
        if n == 0:
            p = a*v.sin() + b*v.cos()
        elif n == 1:
            p = b*v.sin() + c*v.cos()
        elif n == 2:
            p = a*v.sin() + c*v.cos()
        clines = T_est[-1, :3, 4].expand(m_max, 3) + p
        plt.plot(clines[:, i1], clines[:, i2], color=color)
//  [python]
def max_eig_of_hessian(y, x):
    max_eig = torch.eig(hessian(y, x))[0].norm(dim=1).max()

    return max_eig.item()
//  [python]
def max_eig_of_hessian(y, x):
    max_eig = torch.linalg.eig(hessian(y, x))[0].norm(dim=1).max()

    return max_eig.item()
//  [python]
def _matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: matrix
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.eig(matrix, eigenvectors=True)
    # vals, vecs = torch.linalg.eig(matrix)
    vals = torch.view_as_complex(vals.contiguous())+0.001
    vals_pow = vals.pow(p)
    vals_pow = torch.view_as_real(vals_pow)[:, 0]
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs)))
    return matrix_pow
//  [python]
def _matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: matrix
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.linalg.eig(matrix)
    vals = torch.view_as_complex(vals.contiguous())+0.001
    vals_pow = vals.pow(p)
    vals_pow = torch.view_as_real(vals_pow)[:, 0]
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs)))
    return matrix_pow
//  [python]
def forward(self, features):     # features: NCWH
        k = features.size(0) * features.size(2) * features.size(3)
        x_mean = (features.sum(dim=2).sum(dim=2).sum(dim=0) / k).unsqueeze(0).unsqueeze(2).unsqueeze(2)
        features = features - x_mean

        reshaped_features = features.view(features.size(0), features.size(1), -1)\
            .permute(1, 0, 2).contiguous().view(features.size(1), -1)

        cov = torch.matmul(reshaped_features, reshaped_features.t()) / k
        eigval, eigvec = torch.eig(cov, eigenvectors=True)

        first_compo = eigvec[:, 0]

        projected_map = torch.matmul(first_compo.unsqueeze(0), reshaped_features).view(1, features.size(0), -1)\
            .view(features.size(0), features.size(2), features.size(3))

        maxv = projected_map.max()
        minv = projected_map.min()

        projected_map *= (maxv + minv) / torch.abs(maxv + minv)
        
        return projected_map
//  [python]
def forward(self, features):     # features: NCWH
        k = features.size(0) * features.size(2) * features.size(3)
        x_mean = (features.sum(dim=2).sum(dim=2).sum(dim=0) / k).unsqueeze(0).unsqueeze(2).unsqueeze(2)
        features = features - x_mean

        reshaped_features = features.view(features.size(0), features.size(1), -1)\
            .permute(1, 0, 2).contiguous().view(features.size(1), -1)

        cov = torch.matmul(reshaped_features, reshaped_features.t()) / k
        eigval, eigvec = torch.linalg.eig(cov)

        first_compo = eigvec[:, 0]

        projected_map = torch.matmul(first_compo.unsqueeze(0), reshaped_features).view(1, features.size(0), -1)\
            .view(features.size(0), features.size(2), features.size(3))

        maxv = projected_map.max()
        minv = projected_map.min()

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

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

        return data_bamt
//  [python]
def _convolve(dataset: List[Graph], weights: List[float], left_vertices: List[int]) -> List[Graph]:
        new_data = []
        for graph in dataset:
            adj = to_dense_adj(graph.edge_index)
            eigs = torch.eig(adj.reshape(adj.shape[1], adj.shape[2]))
            eigenvectors = torch.real(eigs[1])

            eig = torch.real(eigs[0])  # .T[0].T
            ordered, indices = torch.sort(eig[: graph.num_nodes], descending=True)
            lef = indices[left_vertices]
            zeroed = torch.tensor(list(set(range(len(eig))) - set(lef.tolist())))
            if len(zeroed) > 0:
                eig[zeroed] = 0

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

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

            graph.edge_index, graph.edge_weight = dense_to_sparse(convolved)
            graph.edge_weight = graph.edge_weight
            graph.edge_index = graph.edge_index.type(torch.LongTensor)
            new_data.append(graph)
        return new_data
//  [python]
def _convolve(dataset: List[Graph], weights: List[float], left_vertices: List[int]) -> List[Graph]:
        new_data = []
        for graph in dataset:
            adj = to_dense_adj(graph.edge_index)
            eigs = torch.linalg.eig(adj.reshape(adj.shape[1], adj.shape[2]))
            eigenvectors = torch.real(eigs[1])

            eig = torch.real(eigs[0])  # .T[0].T
            ordered, indices = torch.sort(eig[: graph.num_nodes], descending=True)
            lef = indices[left_vertices]
            zeroed = torch.tensor(list(set(range(len(eig))) - set(lef.tolist())))
            if len(zeroed) > 0:
                eig[zeroed] = 0

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

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

            graph.edge_index, graph.edge_weight = dense_to_sparse(convolved)
            graph.edge_weight = graph.edge_weight
            graph.edge_index = graph.edge_index.type(torch.LongTensor)
            new_data.append(graph)
        return new_data
//  [python]
def _compute_nullv(pixels, alpha, intrinsics, least=4):
        # Construct M matrix and find its null eigenvectors with the least eigenvalues
        # Check equation 7 in paper for more details.
        # pixels (..., N, 2); alpha (..., N, 4); intrinsics (..., 3, 3)
        batch, point = pixels.shape[:-2], pixels.shape[-2]
        u, v = pixels[..., 0], pixels[..., 1]
        fu, u0 = intrinsics[..., 0, 0, None], intrinsics[..., 0, 2, None]
        fv, v0 = intrinsics[..., 1, 1, None], intrinsics[..., 1, 2, None]
        a0, a1, a2, a3 = alpha[..., 0], alpha[..., 1], alpha[..., 2], alpha[..., 3]
        O = torch.zeros_like(a1)
        M = torch.stack([a0 * fu, O, a0 * (u0 - u),
                         a1 * fu, O, a1 * (u0 - u),
                         a2 * fu, O, a2 * (u0 - u),
                         a3 * fu, O, a3 * (u0 - u),
                         O, a0 * fv, a0 * (v0 - v),
                         O, a1 * fv, a1 * (v0 - v),
                         O, a2 * fv, a2 * (v0 - v),
                         O, a3 * fv, a3 * (v0 - v)], dim=-1).view(*batch, point * 2, 12)
        eigenvalues, eigenvectors = torch.symeig(M.mT @ M, eigenvectors=True)
        eigenvalues, eigenvectors = eigenvalues.real, eigenvectors.real
        _, index = eigenvalues.topk(k=least, largest=False, sorted=True)  # (batch, 4)
        index = index.flip(dims=[-1]).unsqueeze(-2).tile((1,) * len(batch) + (12, 1))
        return torch.gather(eigenvectors, dim=-1, index=index).mT
//  [python]
def _compute_nullv(pixels, alpha, intrinsics, least=4):
        # Construct M matrix and find its null eigenvectors with the least eigenvalues
        # Check equation 7 in paper for more details.
        # pixels (..., N, 2); alpha (..., N, 4); intrinsics (..., 3, 3)
        batch, point = pixels.shape[:-2], pixels.shape[-2]
        u, v = pixels[..., 0], pixels[..., 1]
        fu, u0 = intrinsics[..., 0, 0, None], intrinsics[..., 0, 2, None]
        fv, v0 = intrinsics[..., 1, 1, None], intrinsics[..., 1, 2, None]
        a0, a1, a2, a3 = alpha[..., 0], alpha[..., 1], alpha[..., 2], alpha[..., 3]
        O = torch.zeros_like(a1)
        M = torch.stack([a0 * fu, O, a0 * (u0 - u),
                         a1 * fu, O, a1 * (u0 - u),
                         a2 * fu, O, a2 * (u0 - u),
                         a3 * fu, O, a3 * (u0 - u),
                         O, a0 * fv, a0 * (v0 - v),
                         O, a1 * fv, a1 * (v0 - v),
                         O, a2 * fv, a2 * (v0 - v),
                         O, a3 * fv, a3 * (v0 - v)], dim=-1).view(*batch, point * 2, 12)
        eigenvalues, eigenvectors = torch.linalg.eig(M.mT @ M)
        eigenvalues, eigenvectors = eigenvalues.real, eigenvectors.real
        _, index = eigenvalues.topk(k=least, largest=False, sorted=True)  # (batch, 4)
        index = index.flip(dims=[-1]).unsqueeze(-2).tile((1,) * len(batch) + (12, 1))
        return torch.gather(eigenvectors, dim=-1, index=index).mT
//  [python]
def match_color(content, style, eps=1e-5):
    content_pixels, style_pixels = content.view(3, -1), style.view(3, -1)
    mean_content, mean_style = content_pixels.mean(1), style_pixels.mean(1)
    dif_content = content_pixels - mean_content.view(3, 1)
    dif_style = style_pixels - mean_style.view(3, 1)
    eps = eps * torch.eye(3, device=content.device)
    cov_content = dif_content @ dif_content.t() / content_pixels.size(1) + eps
    cov_style = dif_style @ dif_style.t() / style_pixels.size(1) + eps
    eval_content, evec_content = torch.eig(cov_content, eigenvectors=True)
    eval_content = eval_content[:, 0].real.diag()
    evec_content = evec_content.real
    eval_style, evec_style = torch.eig(cov_style, eigenvectors=True)
    eval_style = eval_style[:, 0].real.diag()
    evec_style = evec_style.real
    cov_content_sqrt = evec_content @ eval_content.sqrt() @ evec_content.t()
    cov_style_sqrt = evec_style @ eval_style.sqrt() @ evec_style.t()
    weights = cov_style_sqrt @ cov_content_sqrt.inverse()
    bias = mean_style - weights @ mean_content
    content_pixels = weights @ content_pixels + bias.view(3, 1)
    return content_pixels.clamp(0,1).view_as(content)
//  [python]
def match_color(content, style, eps=1e-5):
    content_pixels, style_pixels = content.view(3, -1), style.view(3, -1)
    mean_content, mean_style = content_pixels.mean(1), style_pixels.mean(1)
    dif_content = content_pixels - mean_content.view(3, 1)
    dif_style = style_pixels - mean_style.view(3, 1)
    eps = eps * torch.eye(3, device=content.device)
    cov_content = dif_content @ dif_content.t() / content_pixels.size(1) + eps
    cov_style = dif_style @ dif_style.t() / style_pixels.size(1) + eps
    eval_content, evec_content = torch.linalg.eig(cov_content)
    eval_content = eval_content.real.diag()
    evec_content = evec_content.real
    eval_style, evec_style = torch.linalg.eig(cov_style)
    eval_style = eval_style.real.diag()
    evec_style = evec_style.real
    cov_content_sqrt = evec_content @ eval_content.sqrt() @ evec_content.t()
    cov_style_sqrt = evec_style @ eval_style.sqrt() @ evec_style.t()
    weights = cov_style_sqrt @ cov_content_sqrt.inverse()
    bias = mean_style - weights @ mean_content
    content_pixels = weights @ content_pixels + bias.view(3, 1)
    return content_pixels.clamp(0,1).view_as(content)
//  [python]
def create_spheric_poses(cameras, n_steps=120):
    center = torch.as_tensor([0.,0.,0.], dtype=cameras.dtype, device=cameras.device)
    cam_center = F.normalize(cameras.mean(0), p=2, dim=-1) * cameras.mean(0).norm(2)
    eigvecs = torch.eig(cameras.T @ cameras).eigenvectors
    rot_axis = F.normalize(eigvecs[:,1].real.float(), p=2, dim=-1)
    up = rot_axis
    rot_dir = torch.cross(rot_axis, cam_center)
    max_angle = (F.normalize(cameras, p=2, dim=-1) * F.normalize(cam_center, p=2, dim=-1)).sum(-1).acos().max()

    all_c2w = []
    for theta in torch.linspace(-max_angle, max_angle, n_steps):
        cam_pos = cam_center * math.cos(theta) + rot_dir * math.sin(theta)
        l = F.normalize(center - cam_pos, p=2, dim=0)
        s = F.normalize(l.cross(up), p=2, dim=0)
        u = F.normalize(s.cross(l), p=2, dim=0)
        c2w = torch.cat([torch.stack([s, u, -l], dim=1), cam_pos[:,None]], axis=1)
        all_c2w.append(c2w)

    all_c2w = torch.stack(all_c2w, dim=0)
    
    return all_c2w
//  [python]
def create_spheric_poses(cameras, n_steps=120):
    center = torch.as_tensor([0.,0.,0.], dtype=cameras.dtype, device=cameras.device)
    cam_center = F.normalize(cameras.mean(0), p=2, dim=-1) * cameras.mean(0).norm(2)
    eigvecs = torch.linalg.eig(cameras.T @ cameras).eigenvectors
    rot_axis = F.normalize(eigvecs[:,1].real.float(), p=2, dim=-1)
    up = rot_axis
    rot_dir = torch.cross(rot_axis, cam_center)
    max_angle = (F.normalize(cameras, p=2, dim=-1) * F.normalize(cam_center, p=2, dim=-1)).sum(-1).acos().max()

    all_c2w = []
    for theta in torch.linspace(-max_angle, max_angle, n_steps):
        cam_pos = cam_center * math.cos(theta) + rot_dir * math.sin(theta)
        l = F.normalize(center - cam_pos, p=2, dim=0)
        s = F.normalize(l.cross(up), p=2, dim=0)
        u = F.normalize(s.cross(l), p=2, dim=0)
        c2w = torch.cat([torch.stack([s, u, -l], dim=1), cam_pos[:,None]], axis=1)
        all_c2w.append(c2w)

    all_c2w = torch.stack(all_c2w, dim=0)
    
    return all_c2w
//  [python]
def logm(A):
    det = torch.det(A)
    mask = ~(det > 0)
    b = mask.sum()
    A[mask] = torch.eye(3).unsqueeze(0).to(A).expand(b, -1, -1)
    eigenvalues, eigenvectors = torch.linalg.eig(A)
    return torch.einsum("bij,bj,bjk->bik", eigenvectors, eigenvalues.log(), torch.linalg.inv(eigenvectors)).real
//  [python]
def logm(A):
    det = torch.det(A)
    mask = ~(det > 0)
    b = mask.sum()
    A[mask] = torch.eye(3).unsqueeze(0).to(A).expand(b, -1, -1)
    eigenvalues, eigenvectors = linalg.eig(A)
    return torch.einsum("bij,bj,bjk->bik", eigenvectors, eigenvalues.log(), torch.linalg.inv(eigenvectors)).real
//  [python]
def get_head_agg_weight(num_users, Vars, Hs, *args, **kwargs):

    device = Hs[0][0].device
    num_cls = Hs[0].shape[0] # number of classes
    d = Hs[0].shape[1] # dimension of feature representation
    avg_weight = []
    for i in range(num_users):
        # ---------------------------------------------------------------------------
        # variance ter
        v = torch.tensor(Vars, device=device)
        # ---------------------------------------------------------------------------
        # bias term
        h_ref = Hs[i]
        dist = torch.zeros((num_users, num_users), device=device)
        for j1, j2 in pairwise(tuple(range(num_users))):
            h_j1 = Hs[j1]
            h_j2 = Hs[j2]
            h = torch.zeros((d, d), device=device)
            for k in range(num_cls):
                h += torch.mm((h_ref[k]-h_j1[k]).reshape(d,1), (h_ref[k]-h_j2[k]).reshape(1,d))
            dj12 = torch.trace(h)
            dist[j1][j2] = dj12
            dist[j2][j1] = dj12

        # QP solver
        p_matrix = torch.diag(v) + dist
        p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
        evals, evecs = torch.eig(torch.tensor(p_matrix), eigenvectors=True)
        
        # for numerical stablity
        p_matrix_new = 0
        p_matrix_new = 0
        for ii in range(num_users):
            if evals[ii,0] >= 0.01:
                p_matrix_new += evals[ii,0]*torch.mm(evecs[:,ii].reshape(num_users,1), evecs[:,ii].reshape(1, num_users))
        p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix)>=0.0) else p_matrix
        
        # solve QP
        alpha = 0
        eps = 1e-3
        if np.all(np.linalg.eigvals(p_matrix)>=0):
            alphav = cvx.Variable(num_users)
            obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
            prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
            prob.solve()
            alpha = alphav.value
            alpha = [(i)*(i>eps) for i in alpha] # zero-out small weights (<eps)
            if i == 0:
                print('({}) Agg Weights of Classifier Head'.format(i+1))
                print(alpha,'\n')
            
        else:
            alpha = None # if no solution for the optimization problem, use local classifier only
        
        avg_weight.append(alpha)

    return avg_weight
//  [python]
def get_head_agg_weight(num_users, Vars, Hs, *args, **kwargs):

    device = Hs[0][0].device
    num_cls = Hs[0].shape[0] # number of classes
    d = Hs[0].shape[1] # dimension of feature representation
    avg_weight = []
    for i in range(num_users):
        # ---------------------------------------------------------------------------
        # variance ter
        v = torch.tensor(Vars, device=device)
        # ---------------------------------------------------------------------------
        # bias term
        h_ref = Hs[i]
        dist = torch.zeros((num_users, num_users), device=device)
        for j1, j2 in pairwise(tuple(range(num_users))):
            h_j1 = Hs[j1]
            h_j2 = Hs[j2]
            h = torch.zeros((d, d), device=device)
            for k in range(num_cls):
                h += torch.mm((h_ref[k]-h_j1[k]).reshape(d,1), (h_ref[k]-h_j2[k]).reshape(1,d))
            dj12 = torch.trace(h)
            dist[j1][j2] = dj12
            dist[j2][j1] = dj12

        # QP solver
        p_matrix = torch.diag(v) + dist
        p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
        evals, evecs = torch.linalg.eig(torch.tensor(p_matrix))
        
        # for numerical stablity
        p_matrix_new = 0
        p_matrix_new = 0
        for ii in range(num_users):
            if evals[ii,0] >= 0.01:
                p_matrix_new += evals[ii,0]*torch.mm(evecs[:,ii].reshape(num_users,1), evecs[:,ii].reshape(1, num_users))
        p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix)>=0.0) else p_matrix
        
        # solve QP
        alpha = 0
        eps = 1e-3
        if np.all(np.linalg.eigvals(p_matrix)>=0):
            alphav = cvx.Variable(num_users)
            obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
            prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
            prob.solve()
            alpha = alphav.value
            alpha = [(i)*(i>eps) for i in alpha] # zero-out small weights (<eps)
            if i == 0:
                print('({}) Agg Weights of Classifier Head'.format(i+1))
                print(alpha,'\n')
            
        else:
            alpha = None # if no solution for the optimization problem, use local classifier only
        
        avg_weight.append(alpha)

    return avg_weight
//  [python]
def sqrtm(A):
    det = torch.det(A)
    mask = ~(det > 0)
    b = mask.sum()
    A[mask] = torch.eye(3).unsqueeze(0).to(A).expand(b, -1, -1)    
    eigenvalues, eigenvectors = torch.linalg.eig(A)
    return torch.einsum("bij,bj,bjk->bik", eigenvectors, eigenvalues.sqrt(), torch.linalg.inv(eigenvectors)).real
//  [python]
def sqrtm(A):
    det = torch.det(A)
    mask = ~(det > 0)
    b = mask.sum()
    A[mask] = torch.eye(3).unsqueeze(0).to(A).expand(b, -1, -1)    
    eigenvalues, eigenvectors = linalg.eig(A)
    return torch.einsum("bij,bj,bjk->bik", eigenvectors, eigenvalues.sqrt(), torch.linalg.inv(eigenvectors)).real
//  [python]
def spectral_radius():

    currentdir = os.path.dirname(os.path.realpath(__file__))
    file = "{}/examples/data/spec_radius_opt_data.mat".format(currentdir)
    mat = scipy.io.loadmat(file)
    mat_struct = mat['sys']
    mat_struct = mat_struct[0,0]
    A = torch.from_numpy(mat_struct['A']).to(device=device, dtype=torch_dtype)
    B = torch.from_numpy(mat_struct['B']).to(device=device, dtype=torch_dtype)
    C = torch.from_numpy(mat_struct['C']).to(device=device, dtype=torch_dtype)
    p = B.shape[1]
    m = C.shape[0]
    stability_margin = 1

    # variables and corresponding dimensions.
    var_in = {"X": [p,m] }

    def comb_fn(X_struct):
        # user defined variable, matrix form. torch tensor
        X = X_struct.X

        # objective function
        M           = A + B@X@C
        [D,_]       = t.eig(M)
        f = torch.max(D.imag)

        # inequality constraint, matrix form
        ci = pygransoStruct()
        ci.c1 = torch.max(D.real) + stability_margin

        # equality constraint
        ce = None

        return [f,ci,ce]

    opts = pygransoStruct()
    opts.maxit = 10
    opts.x0 = torch.zeros(p*m,1).to(device=device, dtype=torch_dtype)
    # print for every 10 iterations. default: 1
    opts.print_frequency = 10
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 2/7 passed (Spectral Radius Optimization)")
//  [python]
def spectral_radius():

    currentdir = os.path.dirname(os.path.realpath(__file__))
    file = "{}/examples/data/spec_radius_opt_data.mat".format(currentdir)
    mat = scipy.io.loadmat(file)
    mat_struct = mat['sys']
    mat_struct = mat_struct[0,0]
    A = torch.from_numpy(mat_struct['A']).to(device=device, dtype=torch_dtype)
    B = torch.from_numpy(mat_struct['B']).to(device=device, dtype=torch_dtype)
    C = torch.from_numpy(mat_struct['C']).to(device=device, dtype=torch_dtype)
    p = B.shape[1]
    m = C.shape[0]
    stability_margin = 1

    # variables and corresponding dimensions.
    var_in = {"X": [p,m] }

    def comb_fn(X_struct):
        # user defined variable, matrix form. torch tensor
        X = X_struct.X

        # objective function
        M           = A + B@X@C
        [D,_]       = LA.eig(M)
        f = torch.max(D.imag)

        # inequality constraint, matrix form
        ci = pygransoStruct()
        ci.c1 = torch.max(D.real) + stability_margin

        # equality constraint
        ce = None

        return [f,ci,ce]

    opts = pygransoStruct()
    opts.maxit = 10
    opts.x0 = torch.zeros(p*m,1).to(device=device, dtype=torch_dtype)
    # print for every 10 iterations. default: 1
    opts.print_frequency = 10
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 2/7 passed (Spectral Radius Optimization)")
//  [python]
def _sqrtm(m: torch.Tensor) -> torch.Tensor:
    val, vec = torch.eig(m)
    return vec @ torch.diag(torch.sqrt(val)) @ vec.inverse()
//  [python]
def _sqrtm(m: torch.Tensor) -> torch.Tensor:
    val, vec = torch.linalg.eig(m)
    return vec @ torch.diag(torch.sqrt(val)) @ vec.inverse()
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.linalg.eig(AP) # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    B = contract('ij, j -> i', V_inv, B.to(V)) # V^* B
    P = contract('ij, ...j -> ...i', V_inv, P.to(V)) # V^* P


    return w, P, B, V
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        # p = torch.zeros(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        C = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2] # Only used in testing
        return w, P, B, C, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.eig(AP) # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    C = initial_C(measure, N, dtype=dtype)
    B = contract('ij, j -> i', V_inv, B.to(V)) # V^* B
    C = contract('ij, j -> i', V_inv, C.to(V)) # V^* C
    P = contract('ij, ...j -> ...i', V_inv, P.to(V)) # V^* P


    return w, P, B, C, V
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        # p = torch.zeros(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        C = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2] # Only used in testing
        return w, P, B, C, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype) # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0] # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.linalg.eig(AP) # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    C = initial_C(measure, N, dtype=dtype)
    B = contract('ij, j -> i', V_inv, B.to(V)) # V^* B
    C = contract('ij, j -> i', V_inv, C.to(V)) # V^* C
    P = contract('ij, ...j -> ...i', V_inv, P.to(V)) # V^* P


    return w, P, B, C, V
//  [python]
def _setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C)) # (H C N)
        if self.L.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.L.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.eig(self.dA)
            V_inv = torch.inverse(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def _setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C)) # (H C N)
        if self.L.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.L.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else: raise NotImplementedError("NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def __call__(self, data):
        pos = data.pos

        if self.max_points > 0 and pos.size(0) > self.max_points:
            perm = torch.randperm(pos.size(0))
            pos = pos[perm[:self.max_points]]

        pos = pos - pos.mean(dim=0, keepdim=True)
        C = torch.matmul(pos.t(), pos)
        e, v = torch.eig(C)  # v[:,j] is j-th eigenvector
        e, v = torch.view_as_real(e), v.real

        if self.sort:
            indices = e[:, 0].argsort(descending=True)
            v = v.t()[indices].t()

        data.pos = torch.matmul(data.pos, v)

        if 'normal' in data:
            data.normal = F.normalize(torch.matmul(data.normal, v))

        return data
//  [python]
def __call__(self, data):
        pos = data.pos

        if self.max_points > 0 and pos.size(0) > self.max_points:
            perm = torch.randperm(pos.size(0))
            pos = pos[perm[:self.max_points]]

        pos = pos - pos.mean(dim=0, keepdim=True)
        C = torch.matmul(pos.t(), pos)
        e, v = torch.linalg.eig(C)  # v[:,j] is j-th eigenvector
        e, v = torch.view_as_real(e), v.real

        if self.sort:
            indices = e[:, 0].argsort(descending=True)
            v = v.t()[indices].t()

        data.pos = torch.matmul(data.pos, v)

        if 'normal' in data:
            data.normal = F.normalize(torch.matmul(data.normal, v))

        return data
//  [python]
def Line(self, points_tsr):  # points_tsr.shape = (n_points, 2)
        xs = points_tsr[:, 0]
        min_x = min(xs)
        max_x = max(xs)
        if abs(min_x - max_x).item() <= self.zero_threshold:  # Vertical line
            return (min_x, torch.tensor(0))
        else:  # Non-vertical line
            ys = points_tsr[:, 1]
            min_y = min(ys)
            max_y = max(ys)
            if abs(min_y - max_y).item() <= self.zero_threshold:  # Horizontal line
                return (min_y, torch.tensor(torch.pi/2))
            else:  # Non-horizontal line
                # | x_i    y_i   -1  | | cos(theta) | = | 0 |
                # | ...    ...   ... | | sin(theta) | = | 0 |
                # | ...    ...   ... | |    rho     |   |...|
                A = torch.zeros((points_tsr.shape[0], 3))
                for row in range(points_tsr.shape[0]):
                    x_i = points_tsr[row][0]
                    y_i = points_tsr[row][1]
                    A[row, 0] = x_i
                    A[row, 1] = y_i
                    A[row, 2] = -1
                # Solution to a system of homogeneous linear equations. Cf. https://stackoverflow.com/questions/1835246/how-to-solve-homogeneous-linear-equations-with-numpy
                e_vals, e_vecs = torch.eig(torch.matmul(A.T, A))
                #print(f"DistortionParametersOptimizer.Line(): torch.matmul(A.T, A) = {torch.matmul(A.T, A)}")
                # Extract the eigenvector (column) associated with the minimum eigenvalue
                #print (f"DistortionParametersOptimizer.Line(): e_vals = {e_vals}")
                z = e_vecs[:, torch.argmin(e_vals[:,0])]
                # Multiply by a factor such that cos^2(theta) + sin^2(theta) = 1
                r2 = z[0] ** 2 + z[1] ** 2
                if abs(r2) < self.zero_threshold:
                    raise ValueError(f"DistortionParametersOptimizer.Line(): z[0]**2 + z[1]**2 ({r2}) < {self.zero_threshold}")
                z = z / torch.sqrt(r2)
                #print(f"DistortionParametersOptimizer.Line(): z = {z}")
                theta = torch.angle(torch.complex(z[0].real, z[1].real))
                rho = z[2].real
                return (rho, theta)
//  [python]
def Line(self, points_tsr):  # points_tsr.shape = (n_points, 2)
        xs = points_tsr[:, 0]
        min_x = min(xs)
        max_x = max(xs)
        if abs(min_x - max_x).item() <= self.zero_threshold:  # Vertical line
            return (min_x, torch.tensor(0))
        else:  # Non-vertical line
            ys = points_tsr[:, 1]
            min_y = min(ys)
            max_y = max(ys)
            if abs(min_y - max_y).item() <= self.zero_threshold:  # Horizontal line
                return (min_y, torch.tensor(torch.pi/2))
            else:  # Non-horizontal line
                # | x_i    y_i   -1  | | cos(theta) | = | 0 |
                # | ...    ...   ... | | sin(theta) | = | 0 |
                # | ...    ...   ... | |    rho     |   |...|
                A = torch.zeros((points_tsr.shape[0], 3))
                for row in range(points_tsr.shape[0]):
                    x_i = points_tsr[row][0]
                    y_i = points_tsr[row][1]
                    A[row, 0] = x_i
                    A[row, 1] = y_i
                    A[row, 2] = -1
                # Solution to a system of homogeneous linear equations. Cf. https://stackoverflow.com/questions/1835246/how-to-solve-homogeneous-linear-equations-with-numpy
                e_vals, e_vecs = torch.linalg.eig(torch.matmul(A.T, A))
                #print(f"DistortionParametersOptimizer.Line(): torch.matmul(A.T, A) = {torch.matmul(A.T, A)}")
                # Extract the eigenvector (column) associated with the minimum eigenvalue
                #print (f"DistortionParametersOptimizer.Line(): e_vals = {e_vals}")
                z = e_vecs[:, torch.argmin(e_vals.real)]
                # Multiply by a factor such that cos^2(theta) + sin^2(theta) = 1
                r2 = z[0] ** 2 + z[1] ** 2
                if abs(r2) < self.zero_threshold:
                    raise ValueError(f"DistortionParametersOptimizer.Line(): z[0]**2 + z[1]**2 ({r2}) < {self.zero_threshold}")
                z = z / torch.sqrt(r2)
                #print(f"DistortionParametersOptimizer.Line(): z = {z}")
                theta = torch.angle(torch.complex(z[0].real, z[1].real))
                rho = z[2].real
                return (rho, theta)
//  [python]
def eigenfunctions(self, what='eig'):
        check_is_fitted(self, ['ny_pts_x_', 'ny_pts_y_', 'U_', 'V_', 'fitted_components_'])
        ny_pts_x, ny_pts_y = self.ny_pts_x_, self.ny_pts_y_
        U, V, r = self.U_, self.V_, self.fitted_components_
        C = V.T @ self.kernel.mmv(ny_pts_x, ny_pts_y, U)

        W_X = V.T @ self.kernel.mmv(ny_pts_x, ny_pts_x, V)
        W_Y = U.T @ self.kernel.mmv(ny_pts_y, ny_pts_y, U)

        # Right eigenfunctions
        edec_right = torch.linalg.eig(C)
        edec_left = torch.linalg.eig(C.H)
        edec_left = Eigendecomposition(edec_left.eigenvalues, edec_left.eigenvectors.conj())

        edec_right_topr = fetch_topr_eigensystem(edec_right, r, what='real')
        edec_left_topr = fetch_topr_eigensystem(edec_left, r, what='real')
        evals = edec_right_topr.eigenvalues.conj()

        norm_left = weighted_norm(edec_left_topr.eigenvectors, to_complex(W_X))
        norm_right = weighted_norm(edec_right_topr.eigenvectors, to_complex(W_Y))

        evec_right = edec_right_topr.eigenvectors / norm_right
        evec_left = edec_left_topr.eigenvectors / norm_left

        def efun_right(new_X):
            k = to_complex(self.kernel.mmv(new_X, ny_pts_x, U))
            return k @ evec_right

        def efun_left(new_X):
            k = to_complex(self.kernel.mmv(new_X, ny_pts_y, V))
            return k @ evec_left / evals

        # Normalization required for Koopman modes
        if what == 'modes':
            left_right_dot = torch.sum(evec_left.conj() * evec_right, 0)
            left_right_norms = 1 / left_right_dot
            return evals, evec_right, left_right_norms, efun_left

        return evals, efun_right, efun_left
//  [python]
def compute_LapEig(g, pos_enc_dim): # input g is a DGL graph
    Adj = g.adj().to_dense() # Adjacency matrix
    Dn = ( g.in_degrees()** -0.5 ).diag() # Inverse and sqrt of degree matrix
    Lap = torch.eye(g.number_of_nodes()) - Dn.matmul(Adj).matmul(Dn) # Laplacian operator
    EigVal, EigVec = torch.symeig(Lap, eigenvectors=True) # Compute full EVD
    EigVal, EigVec = EigVal.real, EigVec.real # make eig real
    EigVec = EigVec[:, EigVal.argsort()] # sort in increasing order of eigenvalues
    EigVec = EigVec[:,1:pos_enc_dim+1] # select the first non-trivial "pos_enc_dim" eigenvector
    return EigVec
//  [python]
def compute_LapEig(g, pos_enc_dim): # input g is a DGL graph
    Adj = g.adj().to_dense() # Adjacency matrix
    Dn = ( g.in_degrees()** -0.5 ).diag() # Inverse and sqrt of degree matrix
    Lap = torch.eye(g.number_of_nodes()) - Dn.matmul(Adj).matmul(Dn) # Laplacian operator
    EigVal, EigVec = torch.linalg.eig(Lap) # Compute full EVD
    EigVal, EigVec = EigVal.real, EigVec.real # make eig real
    EigVec = EigVec[:, EigVal.argsort()] # sort in increasing order of eigenvalues
    EigVec = EigVec[:,1:pos_enc_dim+1] # select the first non-trivial "pos_enc_dim" eigenvector
    return EigVec
//  [python]
def _setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping"""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C))  # (H C N)
        if self.L.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.L.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == "linear":
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2 * self.dC[:, :, : self.N]
        elif mode == "diagonal":
            # Eigendecomposition of the A matrix
            L, V = torch.eig(self.dA)
            V_inv = torch.inverse(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print(
                    "Diagonalization error:",
                    torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA),
                )

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract("h n m, h m -> h n", V_inv, self.dB)
            self.dC = contract("h n m, c h n -> c h m", V, self.dC)

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def _setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping"""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        C = _conj(_r2c(self.C))  # (H C N)
        if self.L.item() == 0:
            dC = C
        else:
            # self.C represents C_tilde
            dA_L = power(self.L.item(), self.dA)
            I = torch.eye(self.dA.size(-1)).to(dA_L)

            dC = torch.linalg.solve(
                I - dA_L.transpose(-1, -2),
                C.unsqueeze(-1),
            ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == "linear":
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2 * self.dC[:, :, : self.N]
        elif mode == "diagonal":
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print(
                    "Diagonalization error:",
                    torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA),
                )

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract("h n m, h m -> h n", V_inv, self.dB)
            self.dC = contract("h n m, c h n -> c h m", V, self.dC)

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def pca_feat(self, X, n_components = 3):
        # x should be c*{any shape}
        # conduct normalization
        X = X/torch.norm(X,dim=0,keepdim=True)
        # fit
        X = X.cuda()
        c, *size = X.shape
        X = X.reshape(c,-1).T
        n, c = X.shape
        mean = torch.mean(X, axis=0)
        X = X - mean
        covariance_matrix = 1 / n * torch.matmul(X.T, X)
        eigenvalues, eigenvectors = torch.symeig(covariance_matrix,eigenvectors=True)
        eigenvalues = eigenvalues.real
        eigenvectors = eigenvectors.real
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        proj_mat = eigenvectors[:, 0:n_components]
        # project
        X = X.matmul(proj_mat).T
        X = X.reshape(tuple([-1] + size))
        return X.cpu()
//  [python]
def pca_feat(self, X, n_components = 3):
        # x should be c*{any shape}
        # conduct normalization
        X = X/torch.norm(X,dim=0,keepdim=True)
        # fit
        X = X.cuda()
        c, *size = X.shape
        X = X.reshape(c,-1).T
        n, c = X.shape
        mean = torch.mean(X, axis=0)
        X = X - mean
        covariance_matrix = 1 / n * torch.matmul(X.T, X)
        eigenvalues, eigenvectors = torch.linalg.eig(covariance_matrix)
        eigenvalues = eigenvalues.real
        eigenvectors = eigenvectors.real
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        proj_mat = eigenvectors[:, 0:n_components]
        # project
        X = X.matmul(proj_mat).T
        X = X.reshape(tuple([-1] + size))
        return X.cpu()
//  [python]
def modeig_forward(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape #batch size,channel depth,dimension
    U, S = th.zeros_like(P, device=P.device), th.zeros(batch_size, channels, n, dtype=P.dtype, device=P.device)

    for i in range(batch_size):
        for j in range(channels):
            if(eig_mode=='eig'):
                #This is for v_pytorch >= 1.9;
                s, U[i, j] = th.eig(P[i,j][None,:])
                S[i, j]    = s[:,0]
            elif(eig_mode=='svd'):
                U[i,j], S[i,j], _ = th.svd(P[i,j])

    S_fn = op.fn(S, param)
    X    = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2,3))
    return X, U, S, S_fn
//  [python]
def modeig_forward(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape #batch size,channel depth,dimension
    U, S = th.zeros_like(P, device=P.device), th.zeros(batch_size, channels, n, dtype=P.dtype, device=P.device)

    for i in range(batch_size):
        for j in range(channels):
            if(eig_mode=='eig'):
                #This is for v_pytorch >= 1.9;
                s, U[i, j] = th.linalg.eig(P[i,j][None,:])
                S[i, j]    = s[:,0]
            elif(eig_mode=='svd'):
                U[i,j], S[i,j], _ = th.svd(P[i,j])

    S_fn = op.fn(S, param)
    X    = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2,3))
    return X, U, S, S_fn
//  [python]
def test_eig_check_magma(self, device, dtype):
        # For CUDA inputs only matrices of size larger than 2048x2048 actually call MAGMA library
        shape = (2049, 2049)
        a = make_tensor(shape, dtype=dtype, device=device)
        w, v = torch.eig(a)
        # check correctness using eigendecomposition identity
        self.assertEqual(a.to(v.dtype) @ v, w * v, atol=1e-3, rtol=1e-3)
//  [python]
def test_eig_check_magma(self, device, dtype):
        # For CUDA inputs only matrices of size larger than 2048x2048 actually call MAGMA library
        shape = (2049, 2049)
        a = make_tensor(shape, dtype=dtype, device=device)
        w, v = torch.linalg.eig(a)
        # check correctness using eigendecomposition identity
        self.assertEqual(a.to(v.dtype) @ v, w * v, atol=1e-3, rtol=1e-3)
//  [python]
def eig(self):
        """
        Returns the eigenvalues (optionally eigenvectors) of the linear map used in matrix multiplication.

        :return: (torch.Tensor) Vector of eigenvalues, optionally a tuple including a matrix of eigenvectors.
        """
        return torch.eig(self.effective_W())