Ensure correctness of the axis arg in LayerNorm module.
Ensure correctness of the axis arg in LayerNorm module.

PiperOrigin-RevId: 306603971
Change-Id: Ibc0f3e990c5d796b10a6c6d57aead0a635ba8452

deepmind/dm-haiku,deepmind/dm-haiku