Honor output axes in fake pmap

[sylvesterkaczmarek]

fake_pmap discards out_axes, so requesting axis 1 can return shape (1, 3) instead of the (3, 1) returned by real pmap.

Forward out_axes to vmap. When fake_parallel_axis=True, remove the synthetic axis from the requested position, including PyTree prefixes and unmapped (None) outputs. Add eager/JIT coverage for output axes 0, 1, and 2, nested outputs, and the synthetic-axis mode.

Validation

• Main suite: 602 passed, 84 skipped, and 85 subtests passed.
• Isolated CPU-device configuration test: 1 passed.
• All 50 fake_test.py tests pass with both one and four CPU devices.
• Six new output-axis regression cases fail on unpatched main.
• Repository Flake8 checks and git diff --check pass.

No GPU or TPU execution was performed.