import torch


def safe_cholesky(A, epsilon=1.0e-8):
    """
    Equivalent of torch.linalg.cholesky that progressively adds
    diagonal jitter to avoid cholesky errors.
    """
    if A.shape == (1, 1):
        return A.sqrt()
    try:
        return torch.linalg.cholesky(A)
    except RuntimeError as e:
        Aprime = A.clone()
        jitter_prev = 0.0
        for i in range(5):
            jitter_new = epsilon * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                return torch.linalg.cholesky(Aprime)
            except RuntimeError:
                continue
        raise e