Layer Norm in Flax NN

lnorm.py

class LayerNorm(nn.Module):
       epsilon: float = 1e-6  # Small value for numerical stability
       use_bias: bool = True
       use_scale: bool = True

       @nn.compact
       def __call__(self, x):
           # Calculate mean and variance along the feature dimension
           mean = jnp.mean(x, axis=-1, keepdims=True)
           variance = jnp.var(x, axis=-1, keepdims=True)

           # Normalize the input
           x = (x - mean) / jnp.sqrt(variance + self.epsilon)

           # Apply scaling and bias (optional)
           if self.use_scale:
               scale = self.param('scale', nn.initializers.ones, (x.shape[-1],))
               x = x * scale
           if self.use_bias:
               bias = self.param('bias', nn.initializers.zeros, (x.shape[-1],))
               x = x + bias

           return x