Example 1
Arguments: [
2,
4
]
Expected result: [
[
true,
true,
true,
false
],
[
true,
true,
true,
true
]
]Queries are at positions 2 and 3: the first sees keys 0..2 and the second sees keys 0..3.
Tokenization and attention · Easy
Causal Masks with Cached Positions
Problem 13 of 50 · Python and NumPy
Problem
Return a query_length × key_length boolean causal mask for cached attention. Queries occupy the final query_length positions of the key sequence. Entry [i,j] is true exactly when j ≤ key_length−query_length+i. An empty query sequence returns [].
def cached_causal_mask(query_length, key_length):Constraints
- 0 ≤ query_length ≤ key_length ≤ 32.
- Indices are zero-based; true means attention is allowed.
- Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.
Examples
Example 2
Arguments: [
3,
3
]
Expected result: [
[
true,
false,
false
],
[
true,
true,
false
],
[
true,
true,
true
]
]No cache exists, so the mask is lower triangular including its diagonal.