Batch Loader Under Test

Task

This data loader feeds a training job. It turns a list of token sequences into padded batches, one list of batches per epoch. Its contract is written down, but its only test is a smoke test. Write tests for the contract and fix what fails. Your tests are graded too: they must pass on the loader and fail on five broken variants of it.

Requirements

1. run(request) answers one request. {op: "batches", examples, batch_size, drop_last, shuffle, seed, epochs} returns {epochs}: a list of epochs, each a list of batches {ids, tokens, mask}. Only examples and batch_size are required.
2. Every example is served exactly once per epoch. A last batch smaller than batch_size is kept unless drop_last is true. batch_size and epochs are positive integers (ValueError containing batch_size or epochs otherwise).
3. ids are positions in examples. tokens holds each example padded with 0 to the longest example in its batch. mask has the same shape as tokens, with 1 on a real token and 0 on padding.
4. Without shuffle, every epoch is in input order. With shuffle, one random.Random(seed) is created for the call and each epoch shuffles a fresh list of ids with it, so epochs differ and the same seed repeats the run.
5. {op: "suite"} runs your contract tests against the loader and returns {result: "pass"} or {result: "fail"}. With mutation set to one of the five names in the file, the tests run against that broken variant and must fail. An unknown mutation raises ValueError containing mutation.
6. A contract test is a function in TESTS that takes load, a function with the signature of load_epochs, and uses assert. Do not change the harness in data/mutations.py and the request harness.