Trace the engine's eval kernel under the mesh and the axis rules

[NuojCheng]

Description

fwd_bwd and update wrap their kernel calls in _sharding_ctx() — the engine's
with jax.set_mesh(self._mesh), nn_partitioning.axis_rules(config.logical_axis_rules).
eval_step was the one step path that did not, so the eval kernel was traced with no mesh
in context and an empty logical axis rule set.

Under shard_mode=auto that only costs partitioning quality — every
maybe_shard_with_logical becomes a no-op and XLA guesses — which is why it went unnoticed.
Under explicit axis types it is fatal. The MaxText layers call jax.sharding.reshard(x, P(...))
(e.g. models/qwen3.py:940), and a bare PartitionSpec outside a mesh context is an error
rather than a no-op:

ValueError: Using PartitionSpec when you are not under a mesh context is not allowed.
Please pass a NamedSharding instance or enter into a mesh context via `jax.set_mesh`. Got P(None,)

So eval_step does not merely evaluate with bad partitioning under shard_mode=explicit —
it raises on the first call.

The context is entered around the call, not around _compile_eval_for_batch, because
jax.jit is lazy: tracing happens at the call. That is the pattern fwd_bwd and update
already follow, and _sharding_ctx's own docstring spells it out.

Pre-existing; not introduced by any open PR. Independent of #5138 and #5143 — no file overlap
in the source hunk beyond the same method, and this branch is cut from main.

Tests

Unit — tests/post_training/unit/maxtext_engine_test.py:
test_eval_step_traces_under_the_mesh_and_axis_rules reads jax.sharding.get_abstract_mesh()
and nn_partitioning.get_axis_rules() from inside the loss fn, covering both eval branches
(compile()d and eager) via subTest. Verified red on main — both subtests fail with
AssertionError: True is not false : the eval kernel was traced with no mesh in context — and
green with the fix. Full file: 51 passed.

JAX_PLATFORMS=cpu python -m pytest tests/post_training/unit/maxtext_engine_test.py -q

(the JAX_PLATFORMS=cpu matters — the suite silently skips without it)

End-to-end — qwen3.5-35b-a3b on a v7-8 (8 devices), ep=8, ring-of-experts + ragged-sort,
GA=8, bf16, seq 1024, adamw, 23 steps, then 4 eval_step calls:

shard_mode	before	after
explicit	eval_step raises ValueError on the first call	4 steps OK, 139.9 ms median
auto	4 steps OK, 1327.4 ms median	4 steps OK, 1327.4 ms median

Training step time and peak HBM are unchanged (4185.3 ms / 61.40 G explicit, 4043.8 ms /
61.39 G auto) — this only touches the eval path. The 9.5x explicit-vs-auto eval gap is the
point of the fix: with the rules live, the eval kernel gets the partitioning the training
kernels already had.

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 wraps the evaluation kernel execution in eval_step within the sharding context (_sharding_ctx) to ensure that tracing occurs with the correct mesh and axis rules active. It also adds a unit test to verify this behavior for both compiled and uncompiled evaluation paths. There are no review comments, so no further feedback is provided.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!