import torch


def add_diagonal(x, val):
    """
    Adds val to the diagonal of x.
    :param x: Matrix, shape (*, D, D).
    :param val: Value to add to diagonal of x.
    :return: x + diag(val).
    """
    assert x.shape[-2] == x.shape[-1], "x must be square."

    d = (torch.ones(x.shape[-2], device=x.device) * val).diag_embed()

    return x + d