import torch


def robust_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """
    Implements a Cholesky decomposition.
    If the Cholesky decomposition fails, it is retried with increasing added jitter on the diagonal.
    :param matrix: Symmetric positive semi-definite matrix to factorize.
    :return: Approximate cholesky factor L such that (approximately) LL^T = matrix
    """
    eps = 1e-5 * matrix.trace() / matrix.shape[-1]
    L = None
    for i in range(10):
        try:
            L = torch.linalg.cholesky(matrix)
            break
        except RuntimeError:
            print('Increasing jitter for Cholesky decomposition', flush=True)
            matrix += eps * torch.eye(matrix.shape[-1], device=matrix.device, dtype=matrix.dtype)
            eps *= 2
    if L is None:
        raise RuntimeError('Could not Cholesky decompose the matrix')
    return L