fix: reshard traced operands in _compat_top_k

[mvanhorn]

Description

_named_sharding uses operand.sharding when it is a NamedSharding and otherwise jax.typeof(operand).sharding. _replicate_reduction_axis clears the PartitionSpec entry for a partitioned reduction axis and applies that layout with the original with_sharding_constraint, falling back to reshard. _compat_top_k does this before the first top_k so the replication is part of the trace; if no partitioned axis is visible and top_k still raises, the same lookup runs once more, and if it still finds nothing the original exception is re-raised. A new CPU unit test drives the installed shim on a four-device explicit mesh for jitted and eager sharded inputs, axis -1 and axis 1, an already-unsharded vocabulary axis, and an oversized k that must surface the original error.

RL post-training dies while the vLLM engine is initializing, during precompile of gather_logprobs. That path calls jax.lax.top_k on log-probabilities sharded along the vocabulary axis, and JAX 0.11 raises ShardingTypeError (for example float32[8@data,128256@model] with axis 1). The shim only looked up a NamedSharding on operand.sharding, and only after top_k had already raised. Under jit the operand is a tracer, so that attribute is not the NamedSharding JAX checks; the layout is on jax.typeof(operand).sharding. The reduction axis stayed partitioned, the replication never entered the trace, and the retry invoked top_k on the same sharded tracer.

Fixes #5440

Tests

Tests for this live in tests/post_training/unit/tunix_adapter_top_k_test.py.

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.
  Not claimed: no review round ran on this change before submission.
• 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.
  Not claimed: no local test run is recorded for src/maxtext/integration/tunix/tunix_adapter.py.
• 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.
  Not claimed: the same reason as the box above applies.

[gemini-code-assist]

Code Review

This pull request refactors _compat_top_k in tunix_adapter.py to correctly replicate a sharded reduction axis when tracing under JIT, and adds corresponding unit tests. Feedback suggests simplifying _compat_top_k by removing a redundant try...except block, as the second call to _replicate_reduction_axis inside the exception handler is deterministic and will always return None if the first call did.

[gemini-code-assist]

The try...except block in _compat_top_k is redundant. If _replicate_reduction_axis returns None on the first call (meaning no sharded reduction axis was detected), it will also return None on the second call inside the except block because the inputs (operand and axis) are unchanged and the function is deterministic. Thus, the except block will always re-raise the original exception. We can simplify the function by removing the try...except block entirely.

Suggested change

	def _compat_top_k(operand, k, axis=-1):
	"""Compat shim around jax.lax.top_k to reshard sharded reduction operands."""
	"""Compat shim around jax.lax.top_k to reshard sharded reduction operands.
	JAX 0.11 rejects a sharded reduction axis. Replicate that axis before the
	first call when it is visible, including on tracers, so the reshard is part
	of the trace. If the pre-check saw no sharded axis and the op still raises,
	run the same lookup once more. A second failure propagates unchanged.
	"""
	replicated = _replicate_reduction_axis(operand, axis)
	if replicated is not None:
	return _orig_top_k(replicated, k, axis=axis)
	try:
	return _orig_top_k(operand, k, axis=axis)
	except Exception: # pylint: disable=broad-exception-caught
	sharding = getattr(operand, "sharding", None)
	if sharding is not None and hasattr(sharding, "spec") and hasattr(sharding, "mesh"): # pylint: disable=line-too-long
	spec = list(sharding.spec)
	idx = axis if axis >= 0 else len(spec) + axis
	if 0 <= idx < len(spec):
	spec[idx] = None
	target_sharding = jax.sharding.NamedSharding(sharding.mesh, jax.sharding.PartitionSpec(*spec)) # pylint: disable=line-too-long
	try:
	operand = _orig_wsc(operand, target_sharding)
	except Exception: # pylint: disable=broad-exception-caught
	operand = jax.sharding.reshard(operand, target_sharding)
	return _orig_top_k(operand, k, axis=axis)
	replicated = _replicate_reduction_axis(operand, axis)
	if replicated is None:
	raise
	return _orig_top_k(replicated, k, axis=axis)
	def _compat_top_k(operand, k, axis=-1):
	"""Compat shim around jax.lax.top_k to reshard sharded reduction operands.
	JAX 0.11 rejects a sharded reduction axis. Replicate that axis before the
	first call when it is visible, including on tracers, so the reshard is part
	of the trace.
	"""
	replicated = _replicate_reduction_axis(operand, axis)
	if replicated is not None:
	return _orig_top_k(replicated, k, axis=axis)
	return _orig_top_k(operand, k, axis=axis)