The autoregressive generation loop

Background

"Autoregressive" means the model's own output is fed back as input. Each step: crop the running sequence to the last block_size tokens, ask a step function for the next token, append it, repeat.

Problem statement

Implement autoregressive_generate(start, step_fn, n, block_size) returning the sequence after appending n tokens.

Input

• start — list of initial token ids (the prompt).
• step_fn — callable mapping a context list (length ≤ block_size) to the next token id.
• n — number of tokens to generate.
• block_size — max context length passed to step_fn.

Output

Returns a list of length len(start) + n: the prompt with n generated tokens appended.

Examples

Example 1 — echo the last token

Input:  start = [5], step_fn = (ctx -> ctx[-1]), n = 3, block_size = 8
Output: [5, 5, 5, 5]

Example 2 — increment the last token

Input:  start = [0], step_fn = (ctx -> ctx[-1] + 1), n = 3, block_size = 8
Output: [0, 1, 2, 3]

Constraints

• Each step: pass ids[-block_size:] to step_fn, append the returned token.
• step_fn never receives more than block_size tokens.
• Return a list of length len(start) + n.

Notes

• Greedy, temperature, top-k, top-p all slot in as different step_fns — the loop itself is identical.