Actually use layer norm epsilon in encoder/decoder (#133)

We forgot to pass it to the sublayers.

Author: mattdangerw

keras_nlp/layers/transformer_decoder.py

@@ -117,9 +117,15 @@ def _build(self, input_shape):
             bias_initializer=self.bias_initializer,
         )
 
-        self._decoder_attention_layernorm = keras.layers.LayerNormalization()
-        self._enc_dec_attention_layernorm = keras.layers.LayerNormalization()
-        self._feedforward_layernorm = keras.layers.LayerNormalization()
+        self._decoder_attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._enc_dec_attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._feedforward_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
 
         self._self_attention_dropout = keras.layers.Dropout(rate=self.dropout)
         self._enc_dec_attentiondropout = keras.layers.Dropout(

keras_nlp/layers/transformer_encoder.py

@@ -104,8 +104,12 @@ def _build(self, input_shape):
             bias_initializer=self.bias_initializer,
         )
 
-        self._attention_layernorm = keras.layers.LayerNormalization()
-        self._feedforward_layernorm = keras.layers.LayerNormalization()
+        self._attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._feedforward_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
 
         self._attention_dropout = keras.layers.Dropout(rate=self.dropout)