import torch


def causal_mask(seq_len, mask_center=False):
    """Causal masking for attention

    Args:
        seq_len : Length of sequence to mask
        mask_center : Mask itself. Defaults to False.

    Returns:
        (seq_len x seq_len) mask tensor
    """
    return torch.tril(torch.ones((seq_len, seq_len)), diagonal=-int(mask_center))