[hsuan-lun-chiang]

Description

This PR fixes a critical crash in the generate_param_only_checkpoint.py utility when checkpoint unrolling is executed with Flax NNX enabled (enable_nnx=true).

The Problem / Root Cause

In Flax NNX, the model state variables are partitioned differently:

• Parameters (nnx.Param) are scanned over config.param_scan_axis (typically axis 1).
• Non-parameters (e.g., layers, internal states, cache variables) are rank-1 arrays scanned over axis 0 (as defined in nnx_decoders.py).

Prior to this fix, the unrolling utility’s slicing helper slice_ith blindly applied jnp.take(x, i, axis=config.param_scan_axis) to every node in the PyTree. For rank-1 non-parameter variables, attempting to take along axis 1 caused the ValueError: axis 1 is out of bounds for array of dimension 1 crash.

The Solution

We updated slice_ith to dynamically resolve the scan axis for each array in the PyTree:

• If the array dimension is 1 (representing a non-parameter scanned state), we take along axis 0.
• Otherwise, we take along config.param_scan_axis (normally axis 1).

This ensures robust and correct dynamic slicing across both parameters and non-parameters during Flax NNX checkpoint unrolling.

---

Tests

We reproduced the failure on the main branch and verified the fix across both CPU and TPU platforms (using a physical v5p-8 TPU slice).

1. Reproduction on main

Running the unrolling command on the main branch with enable_nnx=true reproduced the exact crash from the Airflow DAG logs:

python3 -m maxtext.utils.generate_param_only_checkpoint \
  src/maxtext/configs/base.yml \
  async_checkpointing=false \
  base_output_directory=/tmp \
  load_parameters_path=gs://runner-maxtext-logs/2026-07-02/7b/scanned_ckpt/0/items \
  run_name=unscanned_ckpt \
  model_name=mistral-7b \
  force_unroll=true \
  enable_nnx=true

logs

2. Verification with Fix Branch

Running the identical unrolling command on this branch successfully unrolled the checkpoint and completed with exit code 0:

logs

---

Checklist

• 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).

[codecov]

Codecov Report

❌ Patch coverage is 72.72727% with 15 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/utils/generate_param_only_checkpoint.py	72.72%	8 Missing and 7 partials ⚠️

📢 Thoughts on this report? Let us know!

[xibinliu]

Do we need handle the scan axis here?

We should promote get_scan_axis from a nested local closure to a general module-level function, and reuse it here.

[hsuan-lun-chiang]

We promoted get_scan_axis and fully integrated it into the unrolling logic for both base model parameters ( _possibly_unroll_params_nnx ) and LoRA adapters ( _possibly_unroll_lora_params_nnx ).

[hsuan-lun-chiang]

Hi @xibinliu, here's more details for the LoRA unrolling:

• We integrated the get_scan_axis utility function. It dynamically inspects the actual PartitionSpec of the LoRA adapter parameter to resolve the exact scan axis, removing the fragile hardcoded dependency.
• We replaced python iterators entirely with simultaneous dual-tree mapping ( jax.tree_util.tree_map ).
• Since LoRA parameter annotations are PartitionSpec objects (which JAX normally treats as a tuple container), we specified is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec).