Design a Weighted, Checkpointable Data Batcher

Question

You are training on several named datasets. A registry can open an iterator for a dataset at a supplied item offset. Design a batcher that mixes examples according to positive integer weights and can resume exactly after a restart.

The supplied registry exposes `get_iterator(name, offset=0)`, returning an iterator beginning at that dataset's zero-based item offset. It does not advance any other dataset. Implement the batcher, not the registry.

```python
class DataBatcher:
    def __init__(self, registry, weights: dict[str, int], batch_size: int): ...
    def next_batch(self) -> list[object]: ...
    def state_dict(self) -> dict: ...
    def load_state_dict(self, state: dict) -> None: ...
```

## Requirements

- Every full batch contains exactly `batch_size` examples while data remains available.
- Allocate the integral share first, then distribute leftover slots using a deterministic largest-remainder carry across batches.
- Break ties by the insertion order of `weights` and emit examples grouped in that same dataset order.
- Track the number consumed from each dataset and reopen an iterator directly at that offset when restoring.
- A restored batcher must emit exactly the same future sequence as an uninterrupted batcher.

For exact allocation, let `W` be the sum of weights. For each dataset, reserve `batch_size * weight // W` slots and retain the remainder `batch_size * weight % W`. Carries start at zero. Before each batch, add each remainder to its carry. For each remaining slot, choose the dataset with the greatest current carry (insertion order breaks ties), allocate one slot, then subtract `W` from that dataset's carry. Emit each dataset's items together in insertion order.

For weights `{"a": 3, "b": 2, "c": 1}` and batch size four, the first three count vectors in `(a, b, c)` order are `(2, 1, 1)`, `(2, 2, 0)`, and `(2, 1, 1)`. Save carries as well as offsets so a restart preserves the next allocation.

Assume all datasets contain enough examples for the base implementation. Reject empty weights, non-positive weights, invalid batch sizes, and incompatible checkpoint schemas.

Use `ValueError` for invalid configuration or checkpoints. Weights and batch size must be integers (not booleans). A checkpoint has exactly `offsets` and `remainder_carry`, each mapping every configured dataset to an integer; offsets are nonnegative, while carries may be negative. Reject invalid checkpoints before changing live state. Return and load independent copies. Restore only into a batcher with the same ordered weights and batch size; this two-field schema does not support configuration changes.

## Implementation approach

1. Implement the divisible-weight case.
2. Add fair deterministic handling of non-divisible batch sizes.
3. Add exact checkpoint and restore behavior.

## Design discussion (not graded)

- Choose and defend a policy for dataset exhaustion, changing weights, and distribution across multiple trainers.
- Compare deterministic quotas with sampling each slot in proportion to weight.