mjx: preserve varying-axis metadata across Warp FFI calls

[guptaishaan]

Fixes #3426.

Credit for finding and diagnosing this goes to @danielpmorton, who filed the issue with a
minimal repro and a working fork. This PR builds directly on that patch; he is credited as
co-author on the commit. Happy to close this in favour of a PR from him if he'd prefer to
land it himself.

Root cause

Warp derives its FFI output types in get_jax_output_type() from plain shape/dtype
descriptors, which carry no manual-axis metadata. Those go to jax.ffi.ffi_call, so every
value coming back has an empty varying set.

Under shard_map(..., check_vma=True) that is unsound rather than merely cosmetic: data that
is genuinely per-device is reported as replicated. The call still succeeds and returns
correct numbers — which is what makes it easy to miss. lax.scan is only where it becomes
fatal, because its carry-type equality check is the first thing to compare the two:

The input carry has type float32[2]{V:replicas} but the corresponding output carry
component has type float32[2], so the manual axis types do not match.

Approach

The fix restores the union of the inputs' varying axes onto the outputs — JAX's standard
propagation rule for an opaque op. Over-marking as varying is sound; under-marking is not.
One visible consequence worth calling out: code that declared out_specs=P() expecting
replication will now get an error where it previously got silently-wrong results.

This is applied in jax_callable_variadic_tuple rather than in the vendored Warp tree.
Every MJX Warp call funnels through that wrapper (bvh.py, smooth.py, render.py,
forward.py, collision_driver.py, and the generated shim), so it covers MJX's whole
surface, and a patch to mjx/mujoco/mjx/third_party/warp/ would be reverted on the next
sync. The proper permanent home is upstream Warp, which would also cover jax_kernel; the
TODO in the code marks the removal condition.

Two details in the implementation that are load-bearing and easy to "simplify" into bugs:

• Only the missing axes are cast. _pcast_funcs has no varying -> varying entry, so
  a blanket pcast over all outputs raises Unsupported pcast.
• The axis tuple is sorted. It is derived from a frozenset, and frozenset iteration
  order depends on per-process string hashing. Under multi-controller SPMD each host traces
  independently, so an unsorted tuple can produce divergent jaxprs across hosts.

Tests

ffi_test.py fakes a two-device mesh on CPU via
--xla_force_host_platform_device_count=2, so this runs under ordinary CI instead of
skipping without a GPU. This works because the vendored Warp registers a Host FFI target;
released warp-lang==1.14.0 registers CUDA only.

Test	Before	After
test_single_call_preserves_varying_axes	fail	pass
test_varying_axes_survive_scan_carry (reported symptom)	fail	pass
test_mixed_inputs_take_union_of_varying_axes	fail	pass
test_replicated_input_stays_replicated (control)	pass	pass

The first asserts the mechanism directly (manual_axis_type.varying) rather than relying on
scan's carry check, so it keeps testing this bug even if scan changes. The control
passes in both states, showing the fix doesn't over-apply. I also confirmed the change is a
no-op under eager, jit, and non-sharded scan.

Open questions for reviewers

• JAX version floor. jax.lax.pcast needs JAX >= 0.10, and mjx/pyproject.toml
  declares a bare jax. I used a hasattr guard rather than adding a floor, since MJX
  otherwise uses no JAX >= 0.10 APIs and a floor would constrain all users for a
  multi-device edge case. I did not establish whether the bug manifests on JAX 0.9 under the
  older .vma/pvary spelling — if it does, the guard skips silently there. Happy to
  switch to a version floor if you'd prefer.
• Verified on a CPU-faked mesh only, not on real multi-GPU.

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.