[NNX] Checkpoint restore: error on weights missing from the checkpoint

[ecnal-cienet]

Description

On the pure_nnx resume path, a weight present in the model but absent from the checkpoint is handled unsafely. Now that proactive scan_layers validation has landed (#4304) and the reactive mismatch guards were removed (#4340), _populate_pure_dict_from_partial fills every absent leaf with a default:

• a weight Orbax omits entirely → silently zero-filled (silent accuracy loss);
• a weight Orbax leaves under partial_restore as an unmaterialized jax.ShapeDtypeStruct → propagated as-is, then failing with a cryptic, low-level compile/allocation error deep in the first train_step.

The scan_layers check only validates the scan_layers flag, so neither case is caught for an arbitrary missing weight (a changed emb_dim/mlp_dim, a different number of layers, a dropped or renamed module, etc.).

This folds a descriptive guard into _populate_pure_dict_from_partial, tracking the tree path and whether the leaf sits under an NNX-only rngs/dropout subtree:

• rngs/dropout absent → filled with a deterministic default, as before (a Linen-layout checkpoint never carried it).
• any other absent-or-unmaterialized weight → governed by the new checkpoint_missing_param_policy:
  • "error" (default): raise naming the parameter path and expected shape.
  • "warn": log a WARNING listing the zero-filled path and continue.

"error" is the default because a zero-filled or unmaterialized weight is a silentcorrectness failure. A surviving Shape DtypeStruct is treated the same as an absent leaf, so the previously cryptic first-train_step compile error becomes an explicit, path-named error at restore time.

Tests

• tests/unit/checkpointing_nnx_load_test.py::TestMissingParamPolicy — present weights pass through; rng/dropout absence defaults even under error; a missing weight raises naming the path; a surviving ShapeDtypeStruct is treated as missing; warn zero-fills and logs; default policy is error.
• tests/unit/checkpointing_nnx_missing_param_test.py — end-to-end through the real create_orbax_checkpoint_manager → maybe_save_checkpoint → load_state_if_possible stack: a weight missing from the checkpoint raises naming its path, warn zero-fills and continues, and a matching-arch restore returns concrete arrays.

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!

[xibinliu]

_assert_no_shaped_dtype_struct has been removed with Jacky's change. Please rebase and probably implement the check in _populate_pure_dict_from_partial.