Convert the epsilon to be the same dtype as the variance. Otherwise, in HLO, we upcast the variance if it is initially in bfloat16.
Convert the epsilon to be the same dtype as the variance. Otherwise, in HLO, we upcast the variance if it is initially in bfloat16.

PiperOrigin-RevId: 366235033
Change-Id: I061371637a146dd1d3f69db147981402a8165a51

deepmind/dm-haiku,deepmind/dm-haiku