Normalize each layer in layer_stack tests to reduce error size dependence on `num_layers`.
Normalize each layer in layer_stack tests to reduce error size dependence on `num_layers`.

PiperOrigin-RevId: 451348630

deepmind/dm-haiku,deepmind/dm-haiku