Search

Try a topic

Notebook

Causal Mask as an Additive Tensor

Code note Deep learning

def causal_mask(T, device):
    mask = torch.zeros(T, T, device=device)
    mask = mask.masked_fill(
        torch.triu(torch.ones(T, T, device=device), diagonal=1).bool(),
        float("-inf"),
    )
    return mask  # (T, T), broadcasts to (B, H, T, T)

Why it matters: Additive masks (0.0 / -inf) compose — sum causal + padding and pass one tensor. Multiplicative masks don’t compose cleanly.