Slice the mask buffer to [:T, :T]

Background

nanoGPT registers a fixed lower-triangular mask sized for the maximum context (block_size × block_size), but each forward pass only needs the top-left T × T corner for the current sequence length:

att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))

So a length-16 batch uses a 16×16 mask, never the full 1024×1024.

Problem statement

Implement slice_mask(mask, T) returning the top-left T × T block.

Input

• mask — square array, shape (block_size, block_size).
• T — int, current sequence length (T ≤ block_size).

Output

Returns an np.ndarray of shape (T, T): mask[:T, :T].

Examples

Example 1

Input:  mask = [[1,0,0,0],[1,1,0,0],[1,1,1,0],[1,1,1,1]], T = 2
Output: [[1, 0], [1, 1]]

Constraints

• Return mask[:T, :T].
• Shape (T, T).

Notes

• Slicing (not materialising a fresh mask per call) is what lets generation use a tiny mask when T is small — including T=1 with a KV cache.