Fix explicit-mesh loss sharding assert on the NNX path

[ecnal-cienet]

Problem

Training deepseek3-671b-batchsplit (which sets shard_mode=explicit) with pure_nnx=True crashes at the first train_step compile:

AssertionError: `with_sharding_constraint` acts as an assert when all axes of
mesh are of type Explicit. The array sharding: P(('data', 'fsdp', 'expert'), None)
did not match the sharding provided: P(('data', 'fsdp', 'expert', 'context'), None).

Root cause

loss_fn has two branches. The Linen branch already constrains xent/z_loss through MaxText's sharding.maybe_shard_with_logical, but the NNX branch (taken when pure_nnx=True) was still using raw Flax nn.with_logical_constraint.

The two resolve the logical axes differently:

• nn.with_logical_constraint maps activation_embed_and_logits_batch to ('data','fsdp','expert','context') — it keeps the size-1 context axis.
• The rest of the model shards activations via create_sharding, which calls remove_size_one_mesh_axis and drops context (size 1 here), so logits — and therefore xent — arrive sharded ('data','fsdp','expert').

Under auto axes the constraint is just a hint, so the mismatch was harmless. Under explicit axes with_sharding_constraint becomes a hard assert that the spec matches exactly, so the extra context axis trips it.

Fix

Convert the NNX branch to sharding.maybe_shard_with_logical, mirroring the already-fixed Linen branch. That helper builds the sharding via create_sharding (so size-1 context is dropped, matching the array) and, under an explicit mesh, calls jax.sharding.reshard instead of the asserting constraint. With context of size 1 the reshard is a no-op — same physical layout, just a matching spec.

Also updates tests/unit/train_nnx_test.py: the NNX branch now reads model.mesh and config.debug_sharding, so the tiny test stub gets a single-device mesh and the config stub gets debug_sharding.

Tests

Before Fix (NNX, failed): https://cloudlogging.app.goo.gl/XF3eVFANEPx43XBG7
After Fix (NNX, passed): https://cloudlogging.app.goo.gl/DRsc4TkhtNm5nV929

Unit: tests/unit/train_nnx_test.py passes (7 passed).

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @ecnal-cienet, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @ecnal-cienet, but I was unable to process your request. Please see the logs for more details.