Does this attention need a causal mask?

Background

The encoder reads the whole source at once — its self-attention is bidirectional, no causal mask. Only the decoder's self-attention is causally masked, to stop a position from peeking at future tokens during autoregressive generation. (Cross-attention isn't causally masked either — the source is fully visible.)

Problem statement

Implement needs_causal_mask(component) returning whether that attention sublayer uses a causal mask.

Input

• component — one of "encoder" (encoder self-attention), "decoder" (decoder self-attention), or "cross" (cross-attention).

Output

Returns a bool: True only for "decoder", else False.

Examples

Example 1 — encoder is bidirectional

Input:  component = "encoder"
Output: False

Example 2 — decoder self-attention is masked

Input:  component = "decoder"
Output: True

Constraints

• True iff component == "decoder".
• Encoder self-attention and cross-attention return False.

Notes

• Causal masking is purely a decoder self-attention concern; it's what makes generation autoregressive without leaking the future.