import torch


def conjugate_gradient_block(A, B, x0=None, tol=1e-2, maxit=None, eps=1e-6):
    """
    Solve the linear system A X = B using conjugate gradient, where
     - A is an abstract linear operator implementing an n*n matrix
     - B is a matrix right hand size of size n*s
     - X0 is an initial guess of size n*s

    Essentially, this runs #s classical conjugate gradient algorithms in 'parallel',
    and terminates when the worst of the #s residuals is below tol.
    """
    X = x0
    R = B - A(X)
    P = R
    Rs_old = torch.norm(R, 2., dim=0) ** 2.  
    tol_scale = torch.mean(Rs_old) 
    
    k = 0
    while True:
        k += 1
        AP = A(P)
        alpha = Rs_old / (torch.sum(P * AP, dim=0) + eps)
        X += P * alpha

        if k == maxit:
            break
        
        R -= AP * alpha

        Rs_new = torch.norm(R, 2., dim=0) ** 2.
        res = torch.max(Rs_new)
        if res < (tol ** 2.) * tol_scale:
            break

        P = R + P * (Rs_new / (Rs_old + eps))
        Rs_old = Rs_new

    return X, k