Fix/nnx unroll scan axis

[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 3.63636% with 53 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/utils/generate_param_only_checkpoint.py	3.63%	53 Missing ⚠️

📢 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.

[xibinliu]

params = nnx.split_state(state.model, nnx.Param)
num_params = max_utils.calculate_num_params_from_pytree(params)

[hsuan-lun-chiang]

Done. Refactored the parameter counting in _read_train_checkpoint using nnx.split_state to isolate the nnx.Param subtree.

[xibinliu]

here, we should use the param_only_state:

param_only_state = nnx.split_state(state, nnx.Param)

with this, you will be able to 1. save the trainable params only 2. avoid handling of other variables in the model.

[hsuan-lun-chiang]

Done. Implemented parameter-only state splitting and unrolling.