fix(datapipes): resume-reproducible set_epoch across readers and transforms

[peterdsharpe]

Description

Every set_epoch implementation in the datapipes stack reseeded its generator with initial_seed() + epoch. That idiom compounds: manual_seed updates initial_seed, so the epoch→seed mapping depends on call history. With base seed 42, sequential set_epoch(1..5) produces seeds 43, 45, 48, 52, 57 — while a run resumed from a checkpoint that calls set_epoch(5) directly gets 47. A resumed run therefore silently draws different subsample blocks and augmentations than an uninterrupted run, breaking seed-matched A/B protocols. Within a single uninterrupted run nothing is wrong (still deterministic), which is why this went unnoticed.

Fix: capture the base seed once at set_generator time and reseed with base_seed + epoch. Applied to all sites:

• readers/mesh.py (MeshReader, DomainMeshReader) — cherry-picked byte-identical from Feat(datapipes): Zarr mesh schema and readers — a zarr storage path for mesh-native CAE pipelines #1798 (which applies the same fix to its new zarr mesh readers), so the two PRs merge cleanly in either order
• readers/numpy.py, readers/zarr.py, readers/tensorstore_zarr.py
• transforms/base.py (Transform), transforms/mesh/base.py (MeshTransform) — with a lazy fallback for generators declared by subclasses without going through set_generator

RNG.md and DISTRIBUTIONS.md documented the drifting idiom as the contract; both are updated to the resume-reproducible semantics. Compose only delegates and needed no change.

Behavior note: epoch ≥ 2 now draws different (finally history-independent) blocks/augmentations than before the fix, so bitwise loss-curve comparisons across this change will differ; training statistics are unaffected.

Tests: sequential-vs-direct set_epoch equivalence (block([1,2,3]) == block([3])) pinned for MeshReader, ZarrReader coordinated subsampling, and the MeshDataset transform pipeline. Full test/datapipes/readers/ + test_mesh_augmentations.py suites pass (114 passed).

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

None.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[peterdsharpe]

/ok to test 268aff0

[greptile-apps]

Greptile Summary

This PR fixes a real reproducibility bug in the datapipes stack where set_epoch() reseeded generators with initial_seed() + epoch. Because manual_seed() also updates initial_seed(), each call shifted the baseline, so a run resumed from a checkpoint at epoch N drew different subsample blocks and augmentations than an uninterrupted run that reached epoch N sequentially. The fix captures the base seed once at set_generator() time and always reseeds with base_seed + epoch, making the epoch→seed mapping history-independent.

• Readers (mesh.py, numpy.py, zarr.py, tensorstore_zarr.py): add _subsample_base_seed initialized to None in __init__ and captured in set_generator(); set_epoch() uses that value directly.
• Transforms (transforms/base.py, transforms/mesh/base.py): same fix via set_generator(); a lazy fallback in set_epoch() handles subclasses that declare _generator directly without going through set_generator(), capturing the seed on the first call.
• Tests: three new test_set_epoch_resume_reproducible tests (reader, zarr reader, transform pipeline) assert block([1,2,3]) == block([3]), directly pinning the fixed invariant; documentation and CHANGELOG updated accordingly.

Important Files Changed

Filename	Overview
physicsnemo/datapipes/readers/mesh.py	Adds _subsample_base_seed captured at set_generator() time; set_epoch() uses base_seed + epoch — correct and consistent with the other readers.
physicsnemo/datapipes/readers/numpy.py	Same seed-capture pattern as mesh.py; _subsample_base_seed initialised to None in __init__ and set in set_generator().
physicsnemo/datapipes/readers/zarr.py	Identical fix applied; _subsample_base_seed properly initialised and captured at set_generator() time.
physicsnemo/datapipes/readers/tensorstore_zarr.py	Same fix as other readers; seed captured at set_generator() time, epoch reseeding is now history-independent.
physicsnemo/datapipes/transforms/base.py	Fix applied via set_generator(); includes a lazy fallback in set_epoch() for subclasses that declare _generator directly. _epoch_base_seed is not declared in __init__, unlike the consistent None initialisation used by the reader classes.
physicsnemo/datapipes/transforms/mesh/base.py	Same lazy-fallback pattern as Transform; _epoch_base_seed not declared in __init__ (minor inconsistency with readers).
test/datapipes/readers/test_mesh_readers.py	Adds test_set_epoch_resume_reproducible that directly asserts block([1,2,3]) == block([3]), pinning the fixed invariant.
test/datapipes/readers/test_zarr.py	Equivalent resume-reproducibility test for ZarrReader with coordinated subsampling.
test/datapipes/transforms/test_mesh_augmentations.py	End-to-end test via MeshDataset + RandomScaleMesh confirms transform pipeline is also resume-reproducible.
physicsnemo/datapipes/RNG.md	Documentation updated to reflect new base_seed + epoch contract and resume-reproducibility guarantee.
physicsnemo/datapipes/transforms/mesh/DISTRIBUTIONS.md	Doc updated to match new epoch-seeding semantics.
CHANGELOG.md	Changelog entry clearly describes the bug, root cause, and fix.

Comments Outside Diff (2)

1. physicsnemo/datapipes/transforms/base.py, line 60-62 (link)

   The _epoch_base_seed attribute is never declared in __init__, so it only appears in __dict__ after set_generator() or the first lazy set_epoch() call. All reader classes (MeshReader, ZarrReader, etc.) explicitly initialise their equivalent _subsample_base_seed: int | None = None in __init__, making the state visible from construction. Doing the same here would make the transform state consistent and easier to reason about with tools like vars() or custom state-dict serialisers.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

2. physicsnemo/datapipes/transforms/mesh/base.py, line 44-45 (link)

   Same as Transform: _epoch_base_seed is not declared in __init__, leaving it absent from __dict__ until set_generator() or the first lazy set_epoch() call. Adding it here for consistency with the reader classes and with the proposed Transform change.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[peterdsharpe]

Closing in favor of #1742, which fixes the reader half of this more thoroughly (per-(epoch, index) SeedSequence-derived generators — history-independent and order-independent) and now, as of 31b6e897 pushed there, also covers the transforms half plus the doc updates from this PR. One residual note: #1798 carries a copy of this PR's readers/mesh.py fix (79d25079) that will conflict with #1742 once it merges to main — the plan is to drop that commit from #1798 at its next update, taking #1742's version.