Test: XLA:TPU host offload fails inside nested lax.scan

[NuojCheng]

Not for merge. Self-contained reproducer for an XLA:TPU bug, to attach to a report for the compiler team.

A jax.checkpoint policy built with save_and_offload_only_these_names(..., offload_dst="pinned_host") compiles when the checkpointed body sits inside a single lax.scan -- the ordinary scan-over-layers layout of a homogeneous model. Wrapping that scan in a second lax.scan, which is how a heterogeneous layer cycle is scanned (over blocks, each block looping over its own layers), makes the same policy fail with an internal post-optimization error.

The script covers three cases and needs only JAX and a TPU:

FLAT one scan over 12 layers compiles
NESTED scan over 4 blocks x inner scan over 3 layers fails
TRIP_COUNT_ONE scan over 12 layers x inner scan of length 1 fails

TRIP_COUNT_ONE computes exactly what FLAT computes; the only difference is a loop that runs once and produces nothing extra. Both failures are the same symptom the Qwen3-Next block scan hits, down to the shape of the bad bitcast.

This blocks remat offload policies on every hybrid-attention model. The one workaround available to us -- flattening the inner loop into a Python loop so a single scan level remains -- unrolls the block body and costs more HBM than the offload recovers.

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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.

[gemini-code-assist]

Code Review

This pull request introduces a minimal reproduction script and documentation for an XLA:TPU issue where host-offload rematerialization policies fail when nested inside jax.lax.scan loops. The feedback points out a potential IndexError in the reproduction script's exception handling when an exception has an empty string representation, and provides a safer fallback implementation.

[gemini-code-assist]

If str(e) is empty or contains only whitespace, str(e).strip().splitlines() will return an empty list. Accessing index 0 on an empty list will raise an IndexError, masking the original compilation error. Consider safely retrieving the first line or falling back to a default representation of the exception.

Suggested change

	first_line = str(e).strip().splitlines()[0]
	lines = str(e).strip().splitlines()
	first_line = lines[0] if lines else repr(e)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!