[NNX] Checkpoint: persist and restore rngs/dropout across resumes

[ecnal-cienet]

Description

A pure_nnx run saves in the Linen on-disk layout, which has no place for NNX-only rngs/dropout state, so that state is stripped on save. On resume it was refilled from a base jax.random.key(0) — every preemption reset the dropout/RNG stream even though the dataset iterator was correctly restored, breaking mathematical continuity of the dropout masks across steps.

This persists the stripped rngs/dropout state as a separate nnx_aux composite item and restores it when present:

• Save: extract the rngs/dropout subtree (extract_rng_state) before reshaping to the Linen layout, and write it as an nnx_aux item. The items payload stays byte-identical to today's Linen layout, so NNX↔Linen interchangeability is preserved.
• Restore: read nnx_aux alongside items and merge it back; fall back to the deterministic default when it is absent (Linen-trained or pre-fix checkpoints).
• create_orbax_checkpoint_manager registers the nnx_aux item handler — required, as Orbax rejects an unregistered composite item at save time (empty/unused for Linen runs). Emergency managers keep the default-fill path.

Tests

• tests/unit/checkpointing_nnx_load_test.py — TestExtractRngState, TestDeepMerge, TestRngStateAuxPersistence (helper-level: extract keeps only rng/dropout; round-trip restores the saved counter not the base default; no-aux dir falls back; save adds/omits nnx_aux correctly). 14 passed.
• tests/unit/checkpointing_nnx_roundtrip_test.py — end-to-end through the real create_orbax_checkpoint_manager → maybe_save_checkpoint → load_state_if_possible stack against an on-disk Orbax checkpoint: the dropout rng count survives a save/restore (and nnx_aux/ is written), and a checkpoint without nnx_aux falls back to the base default. 2 passed. (This e2e is what exposed the missing item-handler registration — the mocked unit tests passed without it.)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!