[Linxiushen]

When with_pmap(..., broadcast_args_to_devices=False) receives in_axes=1, a valid batch-major input of shape (3, 2) is rejected on two devices because the wrapper checks shape[0]. The same mismatch affects mixed argument axes and PyTree prefixes, even though native jax.pmap accepts these calls.

Validate every mapped leaf against its selected axis, expand existing axis prefixes over subtrees, and skip unmapped and static arguments before indexing axes. Keep the existing leading-axis error wording. The regression suite covers a collective, mixed axes, leaf/prefix mappings, mapped-size errors, and a trailing static argument with the short in_axes=(0,) form accepted by native JAX.

Validation on Windows/Python 3.12.14, JAX/jaxlib 0.8.0 and NumPy 2.5.4, with two forced CPU virtual devices in one process:

• Exact final new tests on unchanged base 02bfaa78564cdcac0bf498f40d6e39edeebacd36: 6 failures and 1 passing compatibility control. Four failures are valid positive-axis mapping paths; two check Chex's mapped-size diagnostic.
• Full affected native chex/_src/variants_test.py: 131 passed, including all 7 new cases and 124 existing cases; no skips or reported warnings. Accepted outputs are checked against raw jax.pmap and independent NumPy expected arrays.
• Google-configured Pylint with upstream test.sh rules, the CI Flake8 selection, Ruff on both changed files, Pyink on changed line ranges, and git diff --check pass.

This is native library CPU regression evidence. Full repository/build/docs/dependency matrix and physical GPU/distributed execution were not run; successful negative-axis mapping is not added to the supported contract. No CLA or DCO agreement was signed by the agent.

AI assistance: Codex assisted with reproduction, implementation, regression tests and review.