Add NNX model support with Linen/NNX checkpoint interoperability

[ecnal-cienet]

Description

Key Changes

Model Creation (model_creation_utils.py)

• Add create_nnx_model() for NNX model instantiation with proper parameter sharding
• Support loading both Linen and NNX checkpoint formats in NNX models
• Change RNG splitting for dropout to ensure reproducibility
• Add checkpoint format detection and conversion during loading

Training Script (simple_nnx_train.py)

• Minimal NNX training loop for testing and validation
• Synthetic data generation with proper sharding
• Checkpoint saving in NNX format
• Optimizer state handling compatible with optax

Checkpoint Conversion (linen_nnx_converter.py)

• Bidirectional converter between Linen and NNX formats
• Automatic format detection
• Handles params nesting differences (Linen: params.params.decoder vs NNX: params.decoder)
• Handles optimizer state structure differences

Tests

N/A

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

❌ Patch coverage is 0% with 468 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/simple_nnx_train.py	0.00%	273 Missing ⚠️
...xText/utils/ckpt_conversion/linen_nnx_converter.py	0.00%	185 Missing ⚠️
src/MaxText/model_creation_utils.py	0.00%	10 Missing ⚠️

📢 Thoughts on this report? Let us know!