def __call__(self, hidden_states):
        """
        Construct a layernorm module in the LongT5 style; No bias and no subtraction of mean.
        """
        # layer norm should always be calculated in float32
        variance = jnp.power(hidden_states.astype("f4"), 2).mean(axis=-1, keepdims=True)
        hidden_states = hidden_states / jnp.sqrt(variance + self.eps)

        return self.weight * hidden_states

        Construct a layernorm module in the LongT5 style; No bias and no subtraction of mean.
        
dpkd/transformers/src/transformers/models/longt5/modeling_flax_longt5.py