Attention & TransformersMedium
Causal Attention Mask
~6 mincode completion
Implement causal_mask(seq_len) to return a float NumPy array of shape (seq_len, seq_len).
Examples
Single token mask
- Input
- causal_mask(1)
- Output
- [[1]]
Lower-triangular mask for seq_len=3
- Input
- causal_mask(3)
- Output
- [[1, 0, 0], [1, 1, 0], [1, 1, 1]]
Hints
Hint 1
A triangular helper builds this mask directly. Check whether the diagonal is included.
Hint 2
A common slip here: Returning upper-triangular (triu) instead of lower.
Requirements
Include the diagonal (self-attention is allowed).
Return numeric values 1.0 and 0.0.
Use a vectorized NumPy approach.
Constraints
Allowed library: NumPy only
Time limit: 200 ms, Memory: 64 MB
Try similar problems(4)
Where this shows up
~6 min
••••••••••••••••
8 employers weight this skill
4 frontier labs, 3 AI product companies, 1 enterprise vendor. Top match scores 92.
Python
import numpy as np
def causal_mask(seq_len: int) -> np.ndarray:
"""
Create a lower-triangular causal mask of shape (seq_len, seq_len).
Semantics:
M[i, j] = 1.0 if j <= i, else 0.0
"""
# 1) Start from an all-ones matrix of shape (seq_len, seq_len)
# 2) Keep only the lower-triangular part (including diagonal)
# 3) Return the mask as float values
# YOUR CODE HERE
pass