Add Unit Test for NNX Decoder

[ecnal-cienet]

Description

• Add comprehensive unit tests to verify that NNX Decoder has identical parameter tree structures to Linen Decoder across all supported models
• Add comparison script for debugging Linen/NNX tree structure differences

Tests

• Run pytest tests/unit/decoder_tree_structure_test.py --cache-clear -s -v to verify tree structure parity for all supported models
• Verify that existing model training/inference is not affected by the API changes

Running tests for specific models

Run test for a single model

pytest tests/unit/decoder_tree_structure_test.py --cache-clear -s -v -k "llama2-7b" 

Run tests for multiple specific models

pytest tests/unit/decoder_tree_structure_test.py --cache-clear -s -v -k "llama2-7b or gemma2-2b"

Run tests for a model family

pytest tests/unit/decoder_tree_structure_test.py --cache-clear -s -v -k "llama3"

Run tests for all Gemma models

pytest tests/unit/decoder_tree_structure_test.py --cache-clear -s -v -k "gemma"

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 8.43806% with 510 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/compare_linen_nnx_tree.py	0.00%	253 Missing ⚠️
src/MaxText/layers/nnx_decoders.py	15.30%	249 Missing ⚠️
src/MaxText/layers/models.py	28.57%	4 Missing and 1 partial ⚠️
src/MaxText/layers/multi_token_prediction.py	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!