import torch
import warnings


def safe_cholesky(covariance_matrix, jitter=1e-6):
    """Perform a safe cholesky decomposition of the covariance matrix.

    If cholesky decomposition raises Runtime error, it adds jitter to the covariance
    matrix.

    Parameters
    ----------
    covariance_matrix: torch.Tensor.
        Tensor with dimensions batch x dim x dim.
    jitter: float, optional.
        Jitter to add to the covariance matrix.
    """
    try:
        return torch.cholesky(covariance_matrix)
    except RuntimeError:
        dim = covariance_matrix.shape[-1]
        if jitter > 1:
            # When jitter is too big, then there is some numerical issue and this avoids
            # stack overflow.
            warnings.warn("Jitter too big. Maybe some numerical issue somewhere.")
            return torch.eye(dim)
        return safe_cholesky(
            covariance_matrix + jitter * torch.eye(dim), jitter=10 * jitter
        )