import torch


def jitter(K: torch.Tensor, jitter: float = 1.0e-9) -> torch.Tensor:
    """Add jitter to the diagonal of a matrix.

    Args:
        K: The matrix to which to add the jitter. Should be two-dimensional.
        jitter: The amount of jitter to add.

    Returns:
        The jittered matrix.

    """
    assert K.dim() == 2
    return K + torch.diag(
        torch.full((K.size()[0],), jitter, dtype=K.dtype, device=K.device)
    )