feat: configure standalone checkpoint converters with --pure-nnx flag

[RexBearIU]

Description

This PR configures the standalone GPT-OSS checkpoint converter scripts with a --pure-nnx flag (defaulting to True) and completely bypasses any runtime loader translators or bridges.

Key Details:

1. Scanned Checkpoint Converter (convert_gpt_oss_ckpt.py):
   • Added --pure-nnx command line argument (default: True).
   • Dynamically formats the converted scanned weight layout to either use pure-NNX integer-nested layer keys under decoder.layers (when --pure-nnx=True), or legacy Flax Linen string keys under decoder.layers_X (when --pure-nnx=False).
2. Unscanned Checkpoint Converter (convert_gpt_oss_unscanned_ckpt.py):
   • Added --pure-nnx command line argument (default: True).
   • Dynamically formats individual layer variables under nested "layers" with integer keys matching modern NNX scan_layers=False hierarchy (when --pure-nnx=True), or flat string keys (when --pure-nnx=False).
3. Pristine Core Codebase:
   • Completely restored model_creation_utils.py to its original, pristine state with zero runtime bridge translators or mappings.
   • This ensures absolutely zero risk of regression or side effects for all other models (Llama, Gemma, Mistral) and standard training pipelines.

Note

Relation: This PR provides the essential --pure-nnx converter formatting required by the GPTOSS-20B modular testing pipeline in #4556. It should be reviewed and merged first.

Tests

• Verified that standard and non-NNX model configurations continue to build and load with zero regressions.
• Used to successfully convert GPTOSS-20B safetensors checkpoints into pure-NNX checkpoints, which were validated under feat(test): add modular E2E TPU testing scripts for GPTOSS-20B #4556 with standard MaxText training loops.

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!