Fix: Gemma4 checkpoint conversion failure when enable_nnx=True

[chiajunglien]

Description

When converting Gemma4 checkpoints with enable_nnx=True and scan_layers=True, the process fails with a ValueError because the flattened param_map keys mismatch the actual state keys.

Cause of the bug

In layers/nnx_decoders.py, the constructor initializes the scanned block container under self.scanned_blocks. However, in the _apply_gemma4_scanned_blocks method, the main scan was applied to and assigned to self.layers instead of self.scanned_blocks.

This inconsistency caused JAX to dynamically bind the scanned blocks under self.layers at runtime under NNX, producing path keys of the form params-decoder-layers-layers_{layer_idx}. Because the Gemma4 checkpoint conversion's parameter mapping strictly expects Linen's hardcoded prefix params-decoder-scanned_blocks-layers_{layer_idx}, the conversion failed with: ValueError: maxtext_state_dict must be a subset of flattened param_map

Solution

Correct the variable naming inside layers/nnx_decoders.py's _apply_gemma4_scanned_blocks method to use self.scanned_blocks instead of self.layers.

Tests

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=gemma4-26b \
    base_output_directory=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id} \
    use_multimodal=False \
    scan_layers=true \
    hardware=cpu skip_jax_distributed_system=True \
    checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False

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 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!