//  [python]
def _solve_pbcg(
        self,
        rhs: torch.Tensor,
    ):
        hvp = create_hvp_function(
            self.model,
            self.loss,
            self.train_dataloader,
            precompute_grad=self.precompute_grad,
        )

        # The block variant of conjugate gradient is known to suffer from breakdown,
        # due to the possibility of rank deficiency of the iterates of the parameter
        # matrix P^tAP, which destabilizes the direct solver.
        # The paper `Randomized Nyström Preconditioning,
        # Frangella, Zachary and Tropp, Joel A. and Udell, Madeleine,
        # SIAM J. Matrix Anal. Appl., 2023`
        # proposes a simple orthogonalization pre-processing. However, we observed, that
        # this stabilization only worked for double precision. We thus implement
        # a different stabilization strategy described in
        # `A breakdown-free block conjugate gradient method, Ji, Hao and Li, Yaohang,
        # BIT Numerical Mathematics, 2017`

        def mat_mat(x: torch.Tensor):
            return torch.vmap(
                lambda u: hvp(u) + self.hessian_regularization * u,
                in_dims=1,
                randomness="same",
            )(x)

        X = torch.clone(rhs.T)

        R = (rhs - mat_mat(X)).T
        Z = R if self.pre_conditioner is None else self.pre_conditioner.solve(R)
        P, _, _ = torch.svd(Z, some=False)
        active_indices = torch.as_tensor(list(range(X.shape[-1])), dtype=torch.long)

        maxiter = self.maxiter if self.maxiter is not None else len(rhs) * 10
        y_norm = torch.linalg.norm(rhs, dim=1)
        tol = torch.clamp(self.rtol**2 * y_norm, min=self.atol**2)

        # In the case the parameter dimension is smaller than the number of right
        # hand sides, we do not shrink the indices due to resulting wrong
        # dimensionality of the svd decomposition. We consider this an edge case, which
        # does not need optimization
        shrink_finished_indices = rhs.shape[0] <= rhs.shape[1]

        for k in range(maxiter):
            Q = mat_mat(P).T
            p_t_ap = P.T @ Q
            alpha = torch.linalg.solve(p_t_ap, P.T @ R)
            X[:, active_indices] += P @ alpha
            R -= Q @ alpha

            B = torch.linalg.norm(R, dim=0)
            non_finished_indices = torch.nonzero(B > tol)
            num_remaining_indices = non_finished_indices.numel()
            non_finished_indices = non_finished_indices.squeeze()

            if num_remaining_indices == 1:
                non_finished_indices = non_finished_indices.unsqueeze(-1)

            if num_remaining_indices == 0:
                break

            # Reduce problem size by removing finished columns from the iteration
            if shrink_finished_indices:
                active_indices = active_indices[non_finished_indices]
                R = R[:, non_finished_indices]
                P = P[:, non_finished_indices]
                Q = Q[:, non_finished_indices]
                p_t_ap = p_t_ap[:, non_finished_indices][non_finished_indices, :]
                tol = tol[non_finished_indices]

            Z = R if self.pre_conditioner is None else self.pre_conditioner.solve(R)
            beta = -torch.linalg.solve(p_t_ap, Q.T @ Z)
            Z_tmp = Z + P @ beta

            if Z_tmp.ndim == 1:
                Z_tmp = Z_tmp.unsqueeze(-1)

            # Orthogonalization search directions to stabilize the action of
            # (P^tAP)^{-1}
            P, _, _ = torch.svd(Z_tmp, some=False)

        return X.T