import torch


def add_id(df):
    """
    Add the identity to a tensor with the shape of the Jacobian

    Args:
      df (torch.tensor):  batched `(n,m,m)` tensor

    Returns:
      torch.tensor:       :code:`df` plus a batched identity of the right shape
    """
    return df + torch.eye(df.shape[1], device=df.device).reshape(
        (1,) + df.shape[1:]
    ).repeat(df.shape[0], 1, 1)