Sampling Training Batches

The problem

The training corpus is tokenized once, into one long 1-D array of token ids, data. Every training step needs a batch: batch_size random pieces of that array, each with its targets. You will write the function that cuts them.

Each piece is a window of block_size + 1 consecutive tokens, starting at a random offset:

the first block_size tokens are the input x;
the same window shifted by one is the target y.

So at every position the target is the next token: y[b, t] is x[b, t + 1], and the last target is the token right after the window.

Example. data = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], block_size = 3, batch_size = 2, and the random offsets come out as [5, 4].

Row 0, offset 5: x[0] = [15, 16, 17], y[0] = [16, 17, 18].

Row 1, offset 4: x[1] = [14, 15, 16], y[1] = [15, 16, 17].

Row 0 holds three training examples: predict 16 after 15, 17 after 15 16, and 18 after 15 16 17.

The offsets

Draw all the offsets with one call:

ix = rng.integers(
    0,
    len(data) - block_size,
    size=batch_size,
)

The upper bound is exclusive, so the largest offset is len(data) - block_size - 1. With 10 tokens and block_size = 3 that is 6: the window covers positions 6 to 8, and its last target is position 9, the last token. One more and the target would fall off the end.

Building the batch without a loop

Make a table of positions with one row per window: the offset of the row plus 0, 1, ..., block_size - 1. Broadcasting a column of offsets against a row of steps builds it at once.

Example. Offsets [2, 5] and block_size = 3: ix[:, None] is the column [[2], [5]], np.arange(3) is the row [0, 1, 2], and their sum is [[2, 3, 4], [5, 6, 7]]. Indexing data with this table gives x; with the table plus 1, it gives y.

What to implement

get_batch(data, block_size, batch_size, rng): return (x, y), both of shape (batch_size, block_size).

Requirements

Draw the offsets exactly as above, so the tests can reproduce them.
No Python loop over the batch: build the table of positions with broadcasting, then index data once for x and once for y.