Causal Attention Mask

First saw this in pi0. But this basically how you implement masked Self-Attention.

See Masked Attention.