Trainer reuses the same data permutation every epoch

Open Beginner friendly
#9 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
75/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python

Research direction

Locate the trainer code containing the shown epoch loop and inspect how JAX PRNG keys are passed to iterate_dataset. Ensure training receives a freshly split key for each epoch, then verify that successive epochs no longer reuse the same data permutation.

Written by the indexing model from the issue text.

Description

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)
Dominant language
Python
Stars
0
Forks
0
PR merge metrics
No merged PRs in 30d

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from lanl/RMC

All issues in lanl/RMC

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.