Trainer reuses the same data permutation every epoch

[Daniel-B-McBride]

The trainer creates the shuffle key once, outside the epoch loop:

key, subkey1, subkey2 = jax.random.split(key, 3)

for epoch in range(train_epochs + 1):
    for x, y in iterate_dataset(
        train_ds,
        nbatches,
        batch_size,
        subkey1,
        True,
    ):

Since JAX PRNG keys are immutable, subkey1 produces the same permutation every epoch.

Expected behavior is to generate a fresh training key each epoch, e.g.

for epoch in range(train_epochs):
    key, train_key = jax.random.split(key)