Fix PPO input sharding to match pmap axis

[lmj0209]

Summary

Pass PPO's _PMAP_AXIS_NAME to bcast_local_devices when
placing the initial training state.

The new jax.pmap implementation returns arrays with a
NamedSharding whose mesh axis matches the pmap axis name
(i in PPO). The current initial state uses
_device_put_sharded instead. The shapes and devices are
the same, but the sharding objects differ, so the first pmap
result changes the input signature seen by the next host
call.

This change:

• lets bcast_local_devices accept an axis name while
  retaining its existing default;
• reuses that helper in PPO;
• initializes PPO state with the same i axis used by the
  training pmap; and
• adds a focused test for the requested sharding axis.

Validation

• CPU brax.training.pmap_test: passed
• CPU minimal PPO smoke: passed
• Radeon PRO W7900 / gfx1100 / ROCm 7.2.1: input/output
  sharding both use axis i
• Two consecutive ROCm PPO host calls: passed

The default _device_put_sharded value remains unchanged
for other callers.