jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_eigen_AG(self, epsilon: "float") -> "None":
        """Update eigenvalues and eigenvectors of A and G.

        Parameters
        ----------
        epsilon:
            The term added to the denominators to improve numerical stability.

        """
        self._dA, self._QA = torch.symeig(self._A, eigenvectors=True, upper=True)
        self._dG, self._QG = torch.symeig(self._G, eigenvectors=True, upper=True)
        self._dA *= (self._dA > epsilon).type_as(self._dA)
        self._dG *= (self._dG > epsilon).type_as(self._dA)
//  [python]
def update_eigen_AG(self, epsilon: "float") -> "None":
        """Update eigenvalues and eigenvectors of A and G.

        Parameters
        ----------
        epsilon:
            The term added to the denominators to improve numerical stability.

        """
        self._dA, self._QA = torch.linalg.eigh(self._A, UPLO="U")
        self._dG, self._QG = torch.linalg.eigh(self._G, UPLO="U")
        self._dA *= (self._dA > epsilon).type_as(self._dA)
        self._dG *= (self._dG > epsilon).type_as(self._dA)
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    """
    Given a num_init_vecs x num_batch x k x k tridiagonal matrix t_mat,
    returns a num_init_vecs x num_batch x k set of eigenvalues
    and a num_init_vecs x num_batch x k x k set of eigenvectors.

    TODO: make the eigenvalue computations done in batch mode.
    """
    orig_device = t_mat.device
    if settings.verbose_linalg.on():
        settings.verbose_linalg.logger.debug(f"Running symeig on a matrix of size {t_mat.shape}.")

    if t_mat.size(-1) < 32:
        retr = torch.symeig(t_mat.cpu())
    else:
        retr = torch.symeig(t_mat)

    evals, evecs = retr
    mask = evals.ge(0)
    evecs = evecs * mask.type_as(evecs).unsqueeze(-2)
    evals = evals.masked_fill_(~mask, 1)

    return evals.to(orig_device), evecs.to(orig_device)
//  [python]
def lanczos_tridiag_to_diag(t_mat):
    """
    Given a num_init_vecs x num_batch x k x k tridiagonal matrix t_mat,
    returns a num_init_vecs x num_batch x k set of eigenvalues
    and a num_init_vecs x num_batch x k x k set of eigenvectors.

    TODO: make the eigenvalue computations done in batch mode.
    """
    orig_device = t_mat.device
    if settings.verbose_linalg.on():
        settings.verbose_linalg.logger.debug(f"Running symeig on a matrix of size {t_mat.shape}.")

    if t_mat.size(-1) < 32:
        retr = torch.linalg.eigh(t_mat.cpu())
    else:
        retr = torch.linalg.eigh(t_mat)

    evals, evecs = retr
    mask = evals.ge(0)
    evecs = evecs * mask.type_as(evecs).unsqueeze(-2)
    evals = evals.masked_fill_(~mask, 1)

    return evals.to(orig_device), evecs.to(orig_device)
//  [python]
def network_svd(ndim, torch_thread, RR_sum, verbose=1):
    s = time.time()
    torch.set_num_threads(torch_thread)
    if verbose == 1:
        print('All networks loaded. Learning vectors via SVD...\n')
    torch.manual_seed(1)
    try:
        d, V = torch.symeig(torch.Tensor(RR_sum), eigenvectors=True)
        del(RR_sum)
        V = V.numpy()[:, ::-1][:, :ndim]
        d = d.numpy()[::-1][:ndim]
        x = np.diag(np.sqrt(np.sqrt(d))).dot(V.T)
        del(d, V)
    except:
        d, V = eigsh(RR_sum, k=ndim)
        del(RR_sum)
        V = V[:, ::-1]
        d = d[::-1]
        x = np.diag(np.sqrt(np.sqrt(d))).dot(np.transpose(V))

    if verbose == 1:
        print(time.time()-s)
        print('Mashup features obtained.\n')
    return x
//  [python]
def network_svd(ndim, torch_thread, RR_sum, verbose=1):
    s = time.time()
    torch.set_num_threads(torch_thread)
    if verbose == 1:
        print('All networks loaded. Learning vectors via SVD...\n')
    torch.manual_seed(1)
    try:
        d, V = torch.linalg.eigh(torch.Tensor(RR_sum))
        del(RR_sum)
        V = V.numpy()[:, ::-1][:, :ndim]
        d = d.numpy()[::-1][:ndim]
        x = np.diag(np.sqrt(np.sqrt(d))).dot(V.T)
        del(d, V)
    except:
        d, V = eigsh(RR_sum, k=ndim)
        del(RR_sum)
        V = V[:, ::-1]
        d = d[::-1]
        x = np.diag(np.sqrt(np.sqrt(d))).dot(np.transpose(V))

    if verbose == 1:
        print(time.time()-s)
        print('Mashup features obtained.\n')
    return x
//  [python]
def alignment_loss_iso_onegraph(pos_y, pos_x, pos_y_mean, pos_x_mean, num_nodes, total_iso):
        with torch.no_grad():
            total_nodes = pos_y.shape[0]
            y = pos_y - pos_y_mean
            x = pos_x - pos_x_mean
            a = y + x
            b = y - x
            a = a.view(-1, 1, 3)
            b = b.view(-1, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(-1, num_nodes, 16)
            tmpb = torch.mean(tmpb, dim=1).view(-1, 4, 4)
            w, v = torch.symeig(tmpb, eigenvectors=True)
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean.expand(total_iso, -1), rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
            pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
            loss = (pos_y - pos_x).norm(dim=-1, keepdim=True).view(-1, num_nodes,).mean(-1)
            return torch.argmin(loss)
//  [python]
def alignment_loss_iso_onegraph(pos_y, pos_x, pos_y_mean, pos_x_mean, num_nodes, total_iso):
        with torch.no_grad():
            total_nodes = pos_y.shape[0]
            y = pos_y - pos_y_mean
            x = pos_x - pos_x_mean
            a = y + x
            b = y - x
            a = a.view(-1, 1, 3)
            b = b.view(-1, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(-1, num_nodes, 16)
            tmpb = torch.mean(tmpb, dim=1).view(-1, 4, 4)
            w, v = torch.linalg.eigh(tmpb)
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean.expand(total_iso, -1), rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
            pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
            loss = (pos_y - pos_x).norm(dim=-1, keepdim=True).view(-1, num_nodes,).mean(-1)
            return torch.argmin(loss)
//  [python]
def auto_orient_and_center_poses(
    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True
) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        Tuple of the oriented poses and the transform matrix.
    """

    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        w, v = torch.symeig(translation_diff.T @ translation_diff, eigenvectors=True)
        eigvec = torch.flip(v, dims=(-1,))

        if torch.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def auto_orient_and_center_poses(
    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True
) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        Tuple of the oriented poses and the transform matrix.
    """

    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.linalg.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def rot_to_quat(rot: torch.Tensor):
    if rot.shape[-2:] != (3, 3):
        raise ValueError('Input rotation is incorrectly shaped')
    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
    k = [[xx + yy + zz, zy - yz, xz - zx, yx - xy], [zy - yz, xx - yy - zz, xy + yx, xz + zx], [xz - zx, xy + yx, yy - xx - zz, yz + zy], [yx - xy, xz + zx, yz + zy, zz - xx - yy]]
    k = 1.0 / 3.0 * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(rot: torch.Tensor):
    if rot.shape[-2:] != (3, 3):
        raise ValueError('Input rotation is incorrectly shaped')
    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
    k = [[xx + yy + zz, zy - yz, xz - zx, yx - xy], [zy - yz, xx - yy - zz, xy + yx, xz + zx], [xz - zx, xy + yx, yy - xx - zz, yz + zy], [yx - xy, xz + zx, yz + zy, zz - xx - yy]]
    k = 1.0 / 3.0 * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def test_eigh_hermitian_grad(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(dims, uplo):
            x = random_hermitian_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = torch.eig(x, eigenvectors=True)
            (w.sum() + abs(v).sum()).backward()
            self.assertEqual(x.grad, x.grad.conj().transpose(-1, -2))  # Check the gradient is Hermitian

        for dims, uplo in itertools.product([(3, 3), (1, 1, 3, 3)], ["L", "U"]):
            run_test(dims, uplo)
//  [python]
def test_eigh_hermitian_grad(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(dims, uplo):
            x = random_hermitian_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = torch.linalg.eigh(x)
            (w.sum() + abs(v).sum()).backward()
            self.assertEqual(x.grad, x.grad.conj().transpose(-1, -2))  # Check the gradient is Hermitian

        for dims, uplo in itertools.product([(3, 3), (1, 1, 3, 3)], ["L", "U"]):
            run_test(dims, uplo)
//  [python]
def get_normals_from_covariance(covariance_mat):
    eigen_values, eigen_vectors = torch.symeig(covariance_mat,eigenvectors=True)
    min_eigenvalue_index = eigen_values.argsort()[:,0]
    return eigen_vectors[torch.arange(eigen_values.shape[0]), :, min_eigenvalue_index]
//  [python]
def get_normals_from_covariance(covariance_mat):
    eigen_values, eigen_vectors = torch.linalg.eigh(covariance_mat)
    min_eigenvalue_index = eigen_values.argsort()[:,0]
    return eigen_vectors[torch.arange(eigen_values.shape[0]), :, min_eigenvalue_index]
//  [python]
def psd_safe_eigen(K):
	Kprime = K.clone()
	jitter = 0
	jitter_new = None
	while True:
		p, q = torch.symeig(Kprime, eigenvectors=True)
		if (p > 0).all():
			if jitter_new is not None:
				warnings.warn(
					f"K not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
			return p, q
		else:
			if jitter == 0:
				jitter_new = 1e-10
			else:
				jitter_new = jitter * 10
		Kprime.diagonal().add_(jitter_new - jitter)
		jitter = jitter_new
//  [python]
def psd_safe_eigen(K):
	Kprime = K.clone()
	jitter = 0
	jitter_new = None
	while True:
		p, q = torch.linalg.eigh(Kprime)
		if (p > 0).all():
			if jitter_new is not None:
				warnings.warn(
					f"K not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
			return p, q
		else:
			if jitter == 0:
				jitter_new = 1e-10
			else:
				jitter_new = jitter * 10
		Kprime.diagonal().add_(jitter_new - jitter)
		jitter = jitter_new
//  [python]
def eigDecompose(laplacian: torch.Tensor, cuda: bool, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Calculate eigenvalues and eigenvectors from the Laplacian matrix.
    """
    if cuda:
        if device is None:
            device = torch.cuda.current_device()
        laplacian = laplacian.float().to(device)
    else:
        laplacian = laplacian.float().to(torch.device('cpu'))
    lambdas, diffusion_map = torch.linalg.eigh(laplacian)
    return lambdas, diffusion_map
//  [python]
def eigDecompose(laplacian: torch.Tensor, cuda: bool, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Calculate eigenvalues and eigenvectors from the Laplacian matrix.
    """
    if cuda:
        if device is None:
            device = torch.cuda.current_device()
        laplacian = laplacian.float().to(device)
    else:
        laplacian = laplacian.float().to(torch.device('cpu'))
    lambdas, diffusion_map = eigh(laplacian)
    return lambdas, diffusion_map
//  [python]
def auto_orient_and_center_poses(
    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True
) -> TensorType["num_poses":..., 3, 4]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        The oriented poses.
    """

    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        _, eigvec = torch.symeig(translation_diff.T @ translation_diff)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def auto_orient_and_center_poses(
    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True
) -> TensorType["num_poses":..., 3, 4]:
    """Orients and centers the poses. We provide two methods for orientation: pca and up.

    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.


    Args:
        poses: The poses to orient.
        method: The method to use for orientation.
        center_poses: If True, the poses are centered around the origin.

    Returns:
        The oriented poses.
    """

    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation_diff = translation - mean_translation

    if center_poses:
        translation = mean_translation
    else:
        translation = torch.zeros_like(mean_translation)

    if method == "pca":
        _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.linalg.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
        oriented_poses = transform @ poses
    elif method == "none":
        transform = torch.eye(4)
        transform[:3, 3] = -translation
        transform = transform[:3, :]
        oriented_poses = transform @ poses

    return oriented_poses, transform
//  [python]
def affine_to_bcdt(affines):
    k3 = affine_to_k3(affines)
    _, Q = torch.symeig(k3)
    bcd = Q[..., :3, -1]
    return torch.cat((bcd, get_affine_translation(affines)), dim=-1)
//  [python]
def affine_to_bcdt(affines):
    k3 = affine_to_k3(affines)
    _, Q = torch.linalg.eigh(k3)
    bcd = Q[..., :3, -1]
    return torch.cat((bcd, get_affine_translation(affines)), dim=-1)
//  [python]
def _synthesize_exact(self) -> Tuple[Tensor, Tensor]:
        r"""Eigendecomposition of explicitly computed Fisher Information Matrix.

        To be used when the input is small (e.g. less than 70x70 image on cluster or 30x30 on your own machine). This
        method obviates the power iteration and its related algorithms (e.g. Lanczos). This method computes the
        Fisher Information Matrix by explicitly computing the Jacobian of the representation wrt the input.

        Returns
        -------
        eig_vals
            Eigenvalues in decreasing order.
        eig_vecs
            Eigenvectors in 2D tensor, whose cols are eigenvectors (i.e. eigendistortions) corresponding to eigenvalues.
        """

        J = self.compute_jacobian()
        F = J.T @ J
        eig_vals, eig_vecs = torch.symeig(F, eigenvectors=True)
        eig_vecs = eig_vecs.flip(dims=(1,))
        eig_vals = eig_vals.flip(dims=(0,))
        return eig_vals, eig_vecs
//  [python]
def _synthesize_exact(self) -> Tuple[Tensor, Tensor]:
        r"""Eigendecomposition of explicitly computed Fisher Information Matrix.

        To be used when the input is small (e.g. less than 70x70 image on cluster or 30x30 on your own machine). This
        method obviates the power iteration and its related algorithms (e.g. Lanczos). This method computes the
        Fisher Information Matrix by explicitly computing the Jacobian of the representation wrt the input.

        Returns
        -------
        eig_vals
            Eigenvalues in decreasing order.
        eig_vecs
            Eigenvectors in 2D tensor, whose cols are eigenvectors (i.e. eigendistortions) corresponding to eigenvalues.
        """

        J = self.compute_jacobian()
        F = J.T @ J
        eig_vals, eig_vecs = torch.linalg.eigh(F, UPLO="U")
        eig_vecs = eig_vecs.flip(dims=(1,))
        eig_vals = eig_vals.flip(dims=(0,))
        return eig_vals, eig_vecs
//  [python]
def pre_transform(self, instance: Data) -> Data:
        # get spectrum
        n = instance.num_nodes
        L_edge_index, L_values = get_laplacian(instance.edge_index, normalization="sym", num_nodes=n)   # [2, X], [X]
        L = to_dense_adj(L_edge_index, edge_attr=L_values, max_num_nodes=n).squeeze(dim=0)              # [N, N]

        Lambda = torch.zeros(1, self.cfg.pe_dims)   # [1, D_pe]
        V = torch.zeros(n, self.cfg.pe_dims)        # [N, D_pe]

        #d = min(n - 1, self.cfg.pe_dims)   # number of eigen-pairs to use (then we zero-pad up to D_pe)
        d = min(n, self.cfg.pe_dims)   # number of eigen-pairs to use (then we zero-pad up to D_pe)
        eigenvalues, eigenvectors = torch.symeig(L, eigenvectors=True)   # [N], [N, N]
        #Lambda[0, :d] = eigenvalues[1:d + 1]
        #V[:, :d] = eigenvectors[:, 1:d + 1]
        Lambda[0, :d] = eigenvalues[0:d]
        V[:, :d] = eigenvectors[:, 0:d]

        instance.update({"Lambda": Lambda, "V": V})

        return instance
//  [python]
def pre_transform(self, instance: Data) -> Data:
        # get spectrum
        n = instance.num_nodes
        L_edge_index, L_values = get_laplacian(instance.edge_index, normalization="sym", num_nodes=n)   # [2, X], [X]
        L = to_dense_adj(L_edge_index, edge_attr=L_values, max_num_nodes=n).squeeze(dim=0)              # [N, N]

        Lambda = torch.zeros(1, self.cfg.pe_dims)   # [1, D_pe]
        V = torch.zeros(n, self.cfg.pe_dims)        # [N, D_pe]

        #d = min(n - 1, self.cfg.pe_dims)   # number of eigen-pairs to use (then we zero-pad up to D_pe)
        d = min(n, self.cfg.pe_dims)   # number of eigen-pairs to use (then we zero-pad up to D_pe)
        eigenvalues, eigenvectors = torch.linalg.eigh(L)   # [N], [N, N]
        #Lambda[0, :d] = eigenvalues[1:d + 1]
        #V[:, :d] = eigenvectors[:, 1:d + 1]
        Lambda[0, :d] = eigenvalues[0:d]
        V[:, :d] = eigenvectors[:, 0:d]

        instance.update({"Lambda": Lambda, "V": V})

        return instance
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = (r_mat - r_mat.transpose(-1, -2))
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2
    c_angle = (torch.einsum('...ii', r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle)
    scale = (angle / (2 * s_angle))
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[...,0,0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.symeig(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[...,-1,:]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[...,None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = (r_mat - r_mat.transpose(-1, -2))
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2
    c_angle = (torch.einsum('...ii', r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle)
    scale = (angle / (2 * s_angle))
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[...,0,0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.linalg.eigh(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[...,-1,:]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[...,None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def get_whitening_parameters(patches):
    n,c,h,w = patches.shape
    patches_flat = patches.view(n, -1)
    est_patch_covariance = (patches_flat.T @ patches_flat) / n
    eigenvalues, eigenvectors = torch.symeig(est_patch_covariance, eigenvectors=True)
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.T.reshape(c*h*w,c,h,w).flip(0)
//  [python]
def get_whitening_parameters(patches):
    n,c,h,w = patches.shape
    patches_flat = patches.view(n, -1)
    est_patch_covariance = (patches_flat.T @ patches_flat) / n
    eigenvalues, eigenvectors = torch.linalg.eigh(est_patch_covariance, UPLO='U')
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.T.reshape(c*h*w,c,h,w).flip(0)
//  [python]
def _check_point_on_manifold(
        self, x: torch.Tensor, *, atol=1e-5, rtol=1e-5
    ) -> Union[Tuple[bool, Optional[str]], bool]:
        ok = torch.allclose(x, x.transpose(-1, -2), atol=atol, rtol=rtol)
        if not ok:
            return False, "`x != x.transpose` with atol={}, rtol={}".format(atol, rtol)
        e, _ = torch.symeig(x, eigenvectors=True)
        ok = (e > -atol).min()
        if not ok:
            return False, "eigenvalues of x are not all greater than 0."
        return True, None
//  [python]
def _check_point_on_manifold(
        self, x: torch.Tensor, *, atol=1e-5, rtol=1e-5
    ) -> Union[Tuple[bool, Optional[str]], bool]:
        ok = torch.allclose(x, x.transpose(-1, -2), atol=atol, rtol=rtol)
        if not ok:
            return False, "`x != x.transpose` with atol={}, rtol={}".format(atol, rtol)
        e, _ = torch.linalg.eigh(x, "U")
        ok = (e > -atol).min()
        if not ok:
            return False, "eigenvalues of x are not all greater than 0."
        return True, None
//  [python]
def camera_centers_up(cameras):
    cam_centers = torch.cat(
        [cam.get_camera_center() for cam in cameras]
    )

    # fit plane to the camera centers
    cov = torch.cov(cam_centers.transpose(-2, -1))
    # plane_mean = cam_centers.mean(dim=0)
    # cam_centers_c = cam_centers - plane_mean[None]
    # cov = (cam_centers_c.t() @ cam_centers_c) / cam_centers_c.shape[0] # biased
    e_val, e_vec = torch.symeig(cov, eigenvectors=True)
    up = e_vec[:, 0]
    
    # plane_normal = _up_direction(plane_normal, train_dataset)
    camera_up = _up_direction(cameras)

    cos = torch.einsum('ji,i->j', camera_up, up) / (torch.norm(camera_up, dim=1) * torch.norm(up))
    up = up if torch.mean(cos) > 0 else -up
    return up
//  [python]
def camera_centers_up(cameras):
    cam_centers = torch.cat(
        [cam.get_camera_center() for cam in cameras]
    )

    # fit plane to the camera centers
    cov = torch.cov(cam_centers.transpose(-2, -1))
    # plane_mean = cam_centers.mean(dim=0)
    # cam_centers_c = cam_centers - plane_mean[None]
    # cov = (cam_centers_c.t() @ cam_centers_c) / cam_centers_c.shape[0] # biased
    _, e_vec = torch.linalg.eigh(cov)
    up = e_vec[:, 0]
    
    # plane_normal = _up_direction(plane_normal, train_dataset)
    camera_up = _up_direction(cameras)

    cos = torch.einsum('ji,i->j', camera_up, up) / (torch.norm(camera_up, dim=1) * torch.norm(up))
    up = up if torch.mean(cos) > 0 else -up
    return up
//  [python]
def rtf_evd(psd_s: Tensor) -> Tensor:
    r"""Estimate the relative transfer function (RTF) or the steering vector by eigenvalue decomposition.

    .. devices:: CPU CUDA

    .. properties:: TorchScript

    Args:
        psd_s (Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
            Tensor of dimension `(..., freq, channel, channel)`

    Returns:
        Tensor: The estimated complex-valued RTF of target speech.
        Tensor of dimension `(..., freq, channel)`
    """
    if not psd_s.is_complex():
        raise TypeError(f"The type of psd_s must be ``torch.cfloat`` or ``torch.cdouble``. Found {psd_s.dtype}.")
    if psd_s.shape[-1] != psd_s.shape[-2]:
        raise ValueError(f"The last two dimensions of psd_s should be the same. Found {psd_s.shape}.")
    e, v = torch.linalg.eig(psd_s)
    _, idx = torch.max(torch.abs(e), dim=-1)
    rtf = v[torch.arange(v.shape[0]).unsqueeze(-1), idx]
    return rtf
//  [python]
def rtf_evd(psd_s: Tensor) -> Tensor:
    r"""Estimate the relative transfer function (RTF) or the steering vector by eigenvalue decomposition.

    .. devices:: CPU CUDA

    .. properties:: TorchScript

    Args:
        psd_s (Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
            Tensor of dimension `(..., freq, channel, channel)`

    Returns:
        Tensor: The estimated complex-valued RTF of target speech.
        Tensor of dimension `(..., freq, channel)`
    """
    if not psd_s.is_complex():
        raise TypeError(f"The type of psd_s must be ``torch.cfloat`` or ``torch.cdouble``. Found {psd_s.dtype}.")
    if psd_s.shape[-1] != psd_s.shape[-2]:
        raise ValueError(f"The last two dimensions of psd_s should be the same. Found {psd_s.shape}.")
    _, v = torch.linalg.eigh(psd_s)  # v is sorted along with eigenvalues in ascending order
    rtf = v[..., -1]  # choose the eigenvector with max eigenvalue
    return rtf
//  [python]
def _whiten(x, alpha, clip_eigs=1e-9):
    """Compute (partial) whitening transform of x. When alpha=0 it is classic ZCA whitening and columns of x are totally
    decorrelated. When alpha=1, nothing happens.

    Assumes x is already centered.
    """
    e, v = torch.linalg.eigh(x.T @ x / len(x))
    e = torch.clip(e, min=clip_eigs, max=None)
    d = alpha + (1 - alpha) * (e ** -0.5)
    # From right to left, the transformation (1) projects x onto v, (2) divides by stdev in each direction, and (3)
    # rotates back to align with original directions in x-space (ZCA)
    z = v @ torch.diag(d) @ v.T
    # Think of this as (z @ x.T).T, but note z==z.T
    return x @ z
//  [python]
def _whiten(x, alpha, clip_eigs=1e-9):
    """Compute (partial) whitening transform of x. When alpha=0 it is classic ZCA whitening and columns of x are totally
    decorrelated. When alpha=1, nothing happens.

    Assumes x is already centered.
    """
    e, v = eigh(x.T @ x / len(x))
    e = torch.clip(e, min=clip_eigs, max=None)
    d = alpha + (1 - alpha) * (e ** -0.5)
    # From right to left, the transformation (1) projects x onto v, (2) divides by stdev in each direction, and (3)
    # rotates back to align with original directions in x-space (ZCA)
    z = v @ torch.diag(d) @ v.T
    # Think of this as (z @ x.T).T, but note z==z.T
    return x @ z
//  [python]
def quaternion_rotate(
    X: torch.tensor, 
    Y: torch.tensor,
) -> torch.tensor:
    """
    Calculate the rotation

    Args:
        X: (natoms,3) tensor of cartesian coordinates
        Y: (natoms,3) tensor of cartesian coordinates

    Returns:
        rot: Rotation matrix (3,3)
    """
    N = X.size(dim=0)
    W = torch.stack([_makeW(*Y[k]) for k in range(N)])
    Q = torch.stack([_makeQ(*X[k]) for k in range(N)])
    Qt_dot_W = torch.stack([torch.matmul(Q[k].T, W[k]) for k in range(N)])
    A = torch.sum(Qt_dot_W, axis=0)
    eigen = torch.symeig(A)
    r = eigen[1][:, eigen[0].argmax()]
    rot = _quaternion_transform(r)
    return rot
//  [python]
def quaternion_rotate(
    X: torch.tensor, 
    Y: torch.tensor,
) -> torch.tensor:
    """
    Calculate the rotation

    Args:
        X: (natoms,3) tensor of cartesian coordinates
        Y: (natoms,3) tensor of cartesian coordinates

    Returns:
        rot: Rotation matrix (3,3)
    """
    N = X.size(dim=0)
    W = torch.stack([_makeW(*Y[k]) for k in range(N)])
    Q = torch.stack([_makeQ(*X[k]) for k in range(N)])
    Qt_dot_W = torch.stack([torch.matmul(Q[k].T, W[k]) for k in range(N)])
    A = torch.sum(Qt_dot_W, axis=0)
    eigen = torch.linalg.eigh(A)
    r = eigen[1][:, eigen[0].argmax()]
    rot = _quaternion_transform(r)
    return rot
//  [python]
def __call__(self, noisy_data):
        E_t = noisy_data['E_t']
        mask = noisy_data['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]).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

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

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp.type_as(A), batch_eigenvalues.type_as(A)

        elif self.mode == 'all':
            eigvals, eigvectors = torch.symeig(L, eigenvectors=True)
            eigvals = eigvals.type_as(A) / 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=noisy_data['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, noisy_data):
        E_t = noisy_data['E_t']
        mask = noisy_data['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]).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

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

            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)
            return n_connected_comp.type_as(A), batch_eigenvalues.type_as(A)

        elif self.mode == 'all':
            eigvals, eigvectors = torch.linalg.eigh(L)
            eigvals = eigvals.type_as(A) / 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=noisy_data['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 test_symeig_custom():
    X = torch.randn(20, 100)
    M = X @ X.T
    l1, W1 = torch.symeig(M, eigenvectors=True)
    l2, W2 = symeig(M)
    assert torch.allclose(l1, l2)
    assert torch.allclose(W1, W2)
//  [python]
def test_symeig_custom():
    X = torch.randn(20, 100)
    M = X @ X.T
    l1, W1 = torch.linalg.eigh(M, UPLO='U')
    l2, W2 = symeig(M)
    assert torch.allclose(l1, l2)
    assert torch.allclose(W1, W2)
//  [python]
def test_symeig_custom_low_rank():
    X = torch.randn(1000, 10)
    M = X @ X.T
    l1, W1 = torch.symeig(M, upper=True)
    l2, W2 = symeig(M)
    # symeig should fail for low-rank
    assert not torch.all(l1 >= 0.0)
    # test clamping to zeros
    assert torch.all(l2 >= 0.0)
//  [python]
def test_symeig_custom_low_rank():
    X = torch.randn(1000, 10)
    M = X @ X.T
    l1, W1 = torch.linalg.eigh(M, UPLO='U')
    l2, W2 = symeig(M)
    # symeig should fail for low-rank
    assert not torch.all(l1 >= 0.0)
    # test clamping to zeros
    assert torch.all(l2 >= 0.0)
//  [python]
def sqrtm(sigma: Tensor) -> Tensor:
    r"""Returns the square root of a positive semi-definite matrix.

    .. math:: \sqrt{\Sigma} = Q \sqrt{\Lambda} Q^T

    where :math:`Q \Lambda Q^T` is the eigendecomposition of :math:`\Sigma`.

    Args:
        sigma: A positive semi-definite matrix, :math:`(*, D, D)`.

    Example:
        >>> V = torch.randn(4, 4, dtype=torch.double)
        >>> A = V @ V.T
        >>> B = sqrtm(A @ A)
        >>> torch.allclose(A, B)
        True
    """

    L, Q = torch.symeig(sigma)
    L = L.relu().sqrt()

    return Q @ (L[..., None] * Q.mT)
//  [python]
def sqrtm(sigma: Tensor) -> Tensor:
    r"""Returns the square root of a positive semi-definite matrix.

    .. math:: \sqrt{\Sigma} = Q \sqrt{\Lambda} Q^T

    where :math:`Q \Lambda Q^T` is the eigendecomposition of :math:`\Sigma`.

    Args:
        sigma: A positive semi-definite matrix, :math:`(*, D, D)`.

    Example:
        >>> V = torch.randn(4, 4, dtype=torch.double)
        >>> A = V @ V.T
        >>> B = sqrtm(A @ A)
        >>> torch.allclose(A, B)
        True
    """

    L, Q = torch.linalg.eigh(sigma)
    L = L.relu().sqrt()

    return Q @ (L[..., None] * Q.mT)
//  [python]
def append_pca_gmm(self, ax, mean, variance, color=None, alphas=None):
        mean = mean.cpu()
        variance = variance.cpu()
        n_components = mean.size(0)
        if color is None:
            color = 'red'
        if alphas is None:
            alphas = [0.8] * n_components
        else:
            alphas = alphas.cpu().tolist()
        for i in range(n_components):
            mu = mean[i]
            var = variance[i]
            mu_transformed = torch.tensor(self.pca.transform(mu.unsqueeze(0)), dtype=torch.float).squeeze(0)
            eigval, eigvec = torch.symeig(var, eigenvectors=True)
            pca_components = torch.tensor(self.pca.components_, dtype=torch.float)
            eigvec_transformed = torch.matmul(pca_components, eigvec)
            var_projected = torch.matmul(eigvec_transformed, torch.matmul(torch.diag(eigval), eigvec_transformed.T))
            eigval_transformed, eigvec_transformed = torch.symeig(var_projected, eigenvectors=True)
            width = eigval_transformed[0].sqrt() * 2
            height = eigval_transformed[1].sqrt() * 2
            std_scale = 3.0
            angle = torch.atan2(eigvec_transformed[1, 0], eigvec_transformed[0, 0]) * 180 / torch.pi
            ell = Ellipse(mu_transformed, width * std_scale, height * std_scale, angle=angle, fc=to_rgba(color, 0.2 * alphas[i]), ec=to_rgba(color, alphas[i]), lw=3)
            ax.add_artist(ell)
//  [python]
def append_pca_gmm(self, ax, mean, variance, color=None, alphas=None):
        mean = mean.cpu()
        variance = variance.cpu()
        n_components = mean.size(0)
        if color is None:
            color = 'red'
        if alphas is None:
            alphas = [0.8] * n_components
        else:
            alphas = alphas.cpu().tolist()
        for i in range(n_components):
            mu = mean[i]
            var = variance[i]
            mu_transformed = torch.tensor(self.pca.transform(mu.unsqueeze(0)), dtype=torch.float).squeeze(0)
            eigval, eigvec = torch.linalg.eigh(var)
            pca_components = torch.tensor(self.pca.components_, dtype=torch.float)
            eigvec_transformed = torch.matmul(pca_components, eigvec)
            var_projected = torch.matmul(eigvec_transformed, torch.matmul(torch.diag(eigval), eigvec_transformed.T))
            eigval_transformed, eigvec_transformed = torch.linalg.eigh(var_projected)
            width = eigval_transformed[0].sqrt() * 2
            height = eigval_transformed[1].sqrt() * 2
            std_scale = 3.0
            angle = torch.atan2(eigvec_transformed[1, 0], eigvec_transformed[0, 0]) * 180 / torch.pi
            ell = Ellipse(mu_transformed, width * std_scale, height * std_scale, angle=angle, fc=to_rgba(color, 0.2 * alphas[i]), ec=to_rgba(color, alphas[i]), lw=3)
            ax.add_artist(ell)
//  [python]
def proj_least_squares(A, B, reg):
    # print(A.shape)
    # print(B.shape)
    # Compute C = A^T A
    C = torch.matmul(A.transpose(-2, -1), A) + reg * torch.eye(A.shape[-1], device=A.device)

    # Compute the eigenvalues and eigenvectors of C
    eigenvalues, eigenvectors = torch.symeig(C, eigenvectors=True)
    # eigenvalues = torch.maximum( eigenvalues,eigenvalues*0+1e-3  )

    # Diagonal matrix with non-zero eigenvalues in the diagonal
    D_inv = torch.diag_embed(1.0 / torch.maximum(eigenvalues, torch.ones_like(eigenvalues) * 1e-4))

    # Compute the pseudo-inverse of C
    C_inv = torch.matmul(torch.matmul(eigenvectors, D_inv), eigenvectors.transpose(-2, -1))

    # Compute X = C_inv A^T B
    B_proj = torch.matmul(A, torch.matmul(torch.matmul(C_inv, A.transpose(-2, -1)), B))
    return B_proj
//  [python]
def proj_least_squares(A, B, reg):
    # print(A.shape)
    # print(B.shape)
    # Compute C = A^T A
    C = torch.matmul(A.transpose(-2, -1), A) + reg * torch.eye(A.shape[-1], device=A.device)

    # Compute the eigenvalues and eigenvectors of C
    eigenvalues, eigenvectors = torch.linalg.eigh(C)
    # eigenvalues = torch.maximum( eigenvalues,eigenvalues*0+1e-3  )

    # Diagonal matrix with non-zero eigenvalues in the diagonal
    D_inv = torch.diag_embed(1.0 / torch.maximum(eigenvalues, torch.ones_like(eigenvalues) * 1e-4))

    # Compute the pseudo-inverse of C
    C_inv = torch.matmul(torch.matmul(eigenvectors, D_inv), eigenvectors.transpose(-2, -1))

    # Compute X = C_inv A^T B
    B_proj = torch.matmul(A, torch.matmul(torch.matmul(C_inv, A.transpose(-2, -1)), B))
    return B_proj
//  [python]
def form_H_sqrt_and_h(problem, P_sqrt, p):
    """Get parameters of EV(Ax+Bu+c) for commitments"""
    A, B = map(torch.from_numpy, (problem.A_mean, problem.B))

    # form M, H, h
    P = P_sqrt.T @ P_sqrt
    M11 = A.T @ P @ A
    for i in range(P.shape[0]):
        for j in range(P.shape[0]):
            M11[i, j] += torch.trace(P @ torch.from_numpy(problem.cov["A"][i, j]))

    M12 = A.T @ P @ B
    M13 = (A.T @ p).reshape(-1, 1)
    M22 = B.T @ P @ B
    M23 = (B.T @ p).reshape(-1, 1)
    M33 = torch.zeros((1, 1))

    M1 = torch.cat((M11, M12, M13), dim=1)
    M2 = torch.cat((M12.T, M22, M23), dim=1)
    M3 = torch.cat((M13.T, M23.T, M33), dim=0)
    M = torch.cat((M1, M2, M3), dim=0)
    H = M[: problem.n + problem.m, : problem.n + problem.m]
    h = M[-1, :-1]

    # form H_sqrt
    eval, evec = torch.symeig(H)
    H_sqrt = torch.diag(torch.sqrt(torch.clamp(eval, min=0.0, max=None))) @ (evec.T)
    return H_sqrt, h
//  [python]
def form_H_sqrt_and_h(problem, P_sqrt, p):
    """Get parameters of EV(Ax+Bu+c) for commitments"""
    A, B = map(torch.from_numpy, (problem.A_mean, problem.B))

    # form M, H, h
    P = P_sqrt.T @ P_sqrt
    M11 = A.T @ P @ A
    for i in range(P.shape[0]):
        for j in range(P.shape[0]):
            M11[i, j] += torch.trace(P @ torch.from_numpy(problem.cov["A"][i, j]))

    M12 = A.T @ P @ B
    M13 = (A.T @ p).reshape(-1, 1)
    M22 = B.T @ P @ B
    M23 = (B.T @ p).reshape(-1, 1)
    M33 = torch.zeros((1, 1))

    M1 = torch.cat((M11, M12, M13), dim=1)
    M2 = torch.cat((M12.T, M22, M23), dim=1)
    M3 = torch.cat((M13.T, M23.T, M33), dim=1)
    M = torch.cat((M1, M2, M3), dim=0)
    H = M[: problem.n + problem.m, : problem.n + problem.m]
    h = M[-1, :-1]

    # form H_sqrt
    eval, evec = torch.linalg.eigh(H)
    H_sqrt = torch.diag(torch.sqrt(torch.clamp(eval, min=0.0, max=None))) @ (evec.T)
    return H_sqrt, h
//  [python]
def _generalized_eigenvalue_decomposition(a, b):
    cholesky = stable_cholesky(b)
    inv_cholesky = torch.inverse(cholesky)
    # Compute C matrix L⁻1 A L^-T
    cmat = inv_cholesky @ a @ inv_cholesky.conj().transpose(-1, -2)
    # Performing the eigenvalue decomposition
    e_val, e_vec = torch.symeig(cmat,eigenvectors=True)
    # Collecting the eigenvectors
    e_vec = torch.matmul(inv_cholesky.conj().transpose(-1, -2), e_vec)
    return e_val, e_vec
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering."""
    if largest is None:
        largest = False
    E, Z = torch.symeig(A, eigenvectors=True, upper=True)
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z