Fix mask propagation of transformer layers (#139)

Author: chenmoneygithub

keras_nlp/layers/transformer_decoder.py

@@ -94,6 +94,7 @@ def __init__(
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
+        self.supports_masking = True
 
     def _build(self, input_shape):
         # Create layers based on input shape.

keras_nlp/layers/transformer_decoder_test.py

@@ -162,6 +162,18 @@ def test_checkpointing_transformer_decoder(self):
         )
         self.assertAllClose(decoder1_output, decoder2_output)
 
+    def test_mask_propagation(self):
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        encoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        decoder_sequence._keras_mask = mask
+        outputs = decoder(decoder_sequence, encoder_sequence)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_save_model(self):
         encoder_input = keras.Input(shape=[4, 6])
         decoder_input = keras.Input(shape=[4, 6])

keras_nlp/layers/transformer_encoder.py

@@ -89,6 +89,7 @@ def __init__(
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
+        self.supports_masking = True
 
     def _build(self, input_shape):
         # Create layers based on input shape.

keras_nlp/layers/transformer_encoder_test.py

@@ -111,6 +111,17 @@ def test_one_training_step_of_transformer_encoder(self):
         self.assertGreater(len(grad), 1)
         optimizer.apply_gradients(zip(grad, model.trainable_variables))
 
+    def test_mask_propagation(self):
+        encoder = transformer_encoder.TransformerEncoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        inputs = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        inputs._keras_mask = mask
+        outputs = encoder(inputs)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_checkpointing_transformer_encoder(self):
         encoder1 = transformer_encoder.TransformerEncoder(
             intermediate_dim=4,