fix(batch_norm): allow string cross_replica_axis for vmap support

[aditya161205]

Fixes NameError when using hk.BatchNorm with a named jax.vmap axis.

Problem

hk.BatchNorm previously assumed cross_replica_axis was always available (like in pmap). When used with vmap, this caused a NameError: unbound axis name during initialization because the named axis doesn't exist in the init context.

Fix

Added a try-except NameError block around the jax.lax.pmean calls.

• If the axis is unbound (e.g. during init), it now skips the sync and uses local stats.
• Inside vmap, it correctly syncs statistics across the batch dimension.

Verification

• Reproduced the crash locally using the script from BatchNorm vmap support for cross_replica_axis fails tests #812.
• Verified fix: init no longer crashes, and vmap statistics match expected values (std ~1.0).
• Passed existing tests: haiku/_src/batch_norm_test.py

Closes #812
Closes #822

[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.