Repository: keras-team/keras-nlp (python)

Issue:
Add a start and end token packer layer
```
class StartEndPacker(keras.layers.Layer):
    """Adds start and end tokens to a sequence and pads to a fixed length.

    If inputs are batched, input should be a `tf.RaggedTensor`s with shape
    `[batch_size, None]` and will be packed and converted to a dense tensor with
    shape `[batch_size, sequence_length]`.
    If inputs are unbatched, inputs should be dense rank-1 tensors of any shape,
    and will be packed to shape `[sequence_length]`.

    Args:
        sequence_length: The desired output length.
        start_value: The id or token that is to be placed at the start of each
            sequence. The dtype must mach the dtype of the input tensors to
            the layer. If None, no start value will be added.
        end_value: The id or token that is to be placed at the end of each
            input segment. The dtype must mach the dtype of the input tensors to
            the layer. If None, no end value will be added.
        pad_value: The id or token that is to be placed into the unused
            positions after the last segment in the sequence.
    """
    def __init__(
        self,
        sequence_length,
        start_value=None,
        end_value=None,
        pad_value=None,
    )
```

Pull request:
Add StartEndPacker layer
Resolves #220 

Patch:
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index ec1ea1ad..4c2fd2cf 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -17,6 +17,7 @@ from keras_nlp.layers.mlm_head import MLMHead
 from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
+from keras_nlp.layers.start_end_packer import StartEndPacker
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
diff --git a/keras_nlp/layers/start_end_packer.py b/keras_nlp/layers/start_end_packer.py
new file mode 100644
index 00000000..704f4892
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer.py
@@ -0,0 +1,166 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Start End Packer implementation based on `keras.layers.Layer`."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class StartEndPacker(keras.layers.Layer):
+    """Adds start and end tokens to a sequence and pads to a fixed length.
+
+    This layer is useful when tokenizing inputs for tasks like translation,
+    where each sequence should include a start and end marker. It should
+    be called after tokenization. The layer will first trim inputs to fit, then
+    add start/end tokens, and finally pad, if necessary, to `sequence_length`.
+
+    If input is batched, input should be a `tf.RaggedTensor` with shape
+    `[batch_size, None]` and will be packed and converted to a dense tensor with
+    shape `[batch_size, sequence_length]`.
+    If input is unbatched, input should be a dense rank-1 tensor of any shape,
+    and will be packed to shape `[sequence_length]`.
+
+    Args:
+        sequence_length: int. The desired output length.
+        start_value: int/str. The ID or token that is to be placed at the start
+            of each sequence. The dtype must match the dtype of the input
+            tensors to the layer. If None, no start value will be added.
+        end_value: int/str. The ID or token that is to be placed at the end of
+            each input segment. The dtype must match the dtype of the input
+            tensors to the layer. If None, no end value will be added.
+        pad_value: int/str. The ID or token that is to be placed into the
+            unused positions after the last segment in the sequence. If None,
+            0 or "" will be added depending on the dtype of the input tensor.
+
+    Examples:
+
+    Unbatched input (int).
+    >>> input_data = tf.constant([5, 6, 7])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=7, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(7,), dtype=int32, numpy=
+    array([1, 5, 6, 7, 2, 0, 0], dtype=int32)>
+
+    Batched input (int).
+    >>> input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11, 12, 13, 14]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
+    array([[ 1,  5,  6,  7,  2,  0],
+           [ 1,  8,  9, 10, 11,  2]], dtype=int32)>
+
+    Unbatched input (str).
+    >>> input_data = tf.constant(["this", "is", "fun"])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(6,), dtype=string, numpy=
+    array([b'
', b'this', b'is', b'fun', b'
', b'<pad>'], dtype=object)>
+
+    Batched input (str).
+    >>> input_data = tf.ragged.constant([["this", "is", "fun"], ["awesome"]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=string, numpy=
+    array([[b'
', b'this', b'is', b'fun', b'
', b'<pad>'],
+           [b'
', b'awesome', b'
', b'<pad>', b'<pad>', b'<pad>']],
+          dtype=object)>
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value=None,
+        end_value=None,
+        pad_value=None,
+        name=None,
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        self.sequence_length = sequence_length
+        self.start_value = start_value
+        self.end_value = end_value
+        self.pad_value = pad_value
+
+    def call(self, inputs):
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+
+        input_is_dense = isinstance(inputs, tf.Tensor)
+        input_is_ragged = isinstance(inputs, tf.RaggedTensor)
+
+        if input_is_dense:
+            if inputs.shape.rank != 1:
+                raise ValueError(
+                    "Input must either be dense with rank 1, or ragged with "
+                    "rank 2. Received dense input with "
+                    f"rank={inputs.shape.rank}"
+                )
+
+            # Add a new axis at the beginning and convert to ragged tensor.
+            inputs = tf.RaggedTensor.from_tensor(tf.expand_dims(inputs, axis=0))
+        elif input_is_ragged:
+            if inputs.shape.rank != 2:
+                raise ValueError(
+                    "Input must either be dense with rank 1, or ragged with "
+                    "rank 2. Received ragged input with "
+                    f"rank={inputs.shape.rank}"
+                )
+
+        batch_size = tf.shape(inputs)[0]
+
+        # Concatenate start and end tokens.
+        if self.start_value is not None:
+            start_token_id_tensor = tf.fill((batch_size, 1), self.start_value)
+            inputs = tf.concat([start_token_id_tensor, inputs], axis=-1)
+        if self.end_value is not None:
+            end_token_id_tensor = tf.fill((batch_size, 1), self.end_value)
+
+            # Trim to leave room for end token.
+            inputs = inputs[..., : self.sequence_length - 1]
+            inputs = tf.concat([inputs, end_token_id_tensor], axis=-1)
+
+        # Pad to desired length.
+        inputs = inputs.to_tensor(
+            default_value=self.pad_value,
+            shape=(batch_size, self.sequence_length),
+        )
+
+        if input_is_dense:
+            inputs = tf.squeeze(inputs, axis=0)
+
+        return inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self.start_value,
+                "end_value": self.end_value,
+                "pad_value": self.pad_value,
+            }
+        )
+        return config
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index f8cdd105..c71d7283 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -30,12 +30,23 @@ class TransformerDecoder(keras.layers.Layer):
     paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
     can instantiate multiple instances of this class to stack up a decoder.
 
+    This layer will always apply a causal mask to the decoder attention layer.
     This layer will correctly compute an attention mask from an implicit
     Keras padding mask (for example, by passing `mask_zero=True` to a
     `keras.layers.Embedding` layer). See the Masking and Padding
     [guide](https://keras.io/guides/understanding_masking_and_padding/)
     for more details.
 
+    This layer can be called with either one or two inputs. The number of inputs
+    must be consistent across all calls. The options are as follows:
+        `layer(decoder_sequence)`: no cross-attention will be built into the
+            decoder block. This is useful when building a "decoder-only"
+            transformer such as GPT-2.
+        `layer(decoder_sequence, encoder_sequence)`: cross-attention will be
+            built into the decoder block. This is useful when building an
+            "encoder-decoder" transformer, such as the original transformer
+            model described in Attention is All You Need.
+
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
         num_heads: int, the number of heads in MultiHeadAttention.
@@ -102,7 +113,7 @@ class TransformerDecoder(keras.layers.Layer):
         self._built = False
         self.supports_masking = True
 
-    def _build(self, input_shape):
+    def _build(self, input_shape, include_cross_attention):
         # Create layers based on input shape.
         self._built = True
         feature_size = input_shape[-1]
@@ -115,29 +126,37 @@ class TransformerDecoder(keras.layers.Layer):
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
         )
-        self._encoder_decoder_attention_layer = keras.layers.MultiHeadAttention(
-            num_heads=self.num_heads,
-            key_dim=self._attention_head_size,
-            value_dim=feature_size,
-            dropout=self.dropout,
-            kernel_initializer=self.kernel_initializer,
-            bias_initializer=self.bias_initializer,
-        )
 
         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._cross_attention_layer = None
+        if include_cross_attention:
+            # Create layers for cross attention.
+            self._cross_attention_layer = keras.layers.MultiHeadAttention(
+                num_heads=self.num_heads,
+                key_dim=self._attention_head_size,
+                value_dim=feature_size,
+                dropout=self.dropout,
+                kernel_initializer=self.kernel_initializer,
+                bias_initializer=self.bias_initializer,
+            )
+
+            self._cross_attention_layernorm = keras.layers.LayerNormalization(
+                epsilon=self.layer_norm_epsilon,
+            )
+
+            self._cross_attention_dropout = keras.layers.Dropout(
+                rate=self.dropout,
+            )
+
         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(
-            rate=self.dropout,
-        )
+
         # First dense layer in the feedforward network, which maps input
         # feauture size to dimension `self.intermediate_dim`.
         self._intermediate_dense = keras.layers.Dense(
@@ -166,18 +185,20 @@ class TransformerDecoder(keras.layers.Layer):
     def call(
         self,
         decoder_sequence,
-        encoder_sequence,
+        encoder_sequence=None,
         decoder_padding_mask=None,
         decoder_attention_mask=None,
         encoder_padding_mask=None,
         encoder_attention_mask=None,
-        use_causal_mask=False,
     ):
         """Forward pass of the TransformerDecoder.
 
         Args:
             decoder_sequence: a Tensor. The decoder input sequence.
-            encoder_sequence: a Tensor. The decoder input sequence.
+            encoder_sequence: a Tensor. The encoder input sequence. For decoder
+                only models (like GPT2), this should be left None. Once the
+                model is called once without an encoder_sequence, you cannot
+                call it again with encoder_sequence.
             decoder_padding_mask: a boolean Tensor, the padding mask of decoder
                 sequence, must of shape [batch_size, decoder_sequence_length].
             decoder_attention_mask: a boolean Tensor. Customized decoder
@@ -188,29 +209,45 @@ class TransformerDecoder(keras.layers.Layer):
             encoder_attention_mask: a boolean Tensor. Customized encoder
                 sequence mask, must of shape
                 [batch_size, encoder_sequence_length, encoder_sequence_length].
-            use_causal_mask: bool, defaults to False. If true, causal mask
-                (masking out future input) is applied on the decoder sequence.
-
         Returns:
             A Tensor of the same shape as the `decoder_sequence`.
         """
+        has_encoder_sequence = encoder_sequence is not None
         if not self._built:
-            self._build(decoder_sequence.shape)
-        encoder_mask = merge_padding_and_attention_mask(
-            encoder_sequence, encoder_padding_mask, encoder_attention_mask
-        )
+            self._build(decoder_sequence.shape, has_encoder_sequence)
+
+        is_cross_attention = self._cross_attention_layer is not None
+        if not is_cross_attention and has_encoder_sequence:
+            raise ValueError(
+                "The number of call arguments to "
+                "`keras_nlp.layers.TransformerDecoder` should not change. "
+                "Use `layer(decoder_sequence, encoder_sequence)` to "
+                "build a layer with cross attention, or "
+                "`layer(decoder_sequence)` to build a layer without. "
+                "This layer has been built without cross attention, but "
+                "you are trying to call it with encoder_sequence."
+            )
+        elif is_cross_attention and not has_encoder_sequence:
+            raise ValueError(
+                "The number of call arguments to "
+                "`keras_nlp.layers.TransformerDecoder` should not change. "
+                "Use `layer(decoder_sequence, encoder_sequence)` to "
+                "build a layer with cross attention, or "
+                "`layer(decoder_sequence)` to build a layer without. "
+                "This layer has been built with cross attention, but "
+                "you did not provide encoder_sequence."
+            )
         decoder_mask = merge_padding_and_attention_mask(
             decoder_sequence, decoder_padding_mask, decoder_attention_mask
         )
-        if use_causal_mask:
-            causal_mask = tf.cast(
-                compute_causal_mask(decoder_sequence),
-                dtype=tf.int32,
-            )
-            if decoder_mask is None:
-                decoder_mask = causal_mask
-            else:
-                decoder_mask = tf.minimum(decoder_mask, causal_mask)
+        causal_mask = tf.cast(
+            compute_causal_mask(decoder_sequence),
+            dtype=tf.int32,
+        )
+        if decoder_mask is None:
+            decoder_mask = causal_mask
+        else:
+            decoder_mask = tf.minimum(decoder_mask, causal_mask)
 
         # Decoder input self-attention.
         self_attended = self._self_attention_layer(
@@ -220,28 +257,34 @@ class TransformerDecoder(keras.layers.Layer):
             attention_mask=decoder_mask,
         )
         self_attended = self._self_attention_dropout(self_attended)
-        self_attended = self._add_and_norm(
+        attention_output = self._add_and_norm(
             self_attended, decoder_sequence, self._decoder_attention_layernorm
         )
-        # Encoder-decoder attention.
-        encoder_decoder_attended = self._encoder_decoder_attention_layer(
-            query=self_attended,
-            value=encoder_sequence,
-            key=encoder_sequence,
-            attention_mask=encoder_mask,
-        )
-        encoder_decoder_attended = self._enc_dec_attentiondropout(
-            encoder_decoder_attended,
-        )
-        encoder_decoder_attended = self._add_and_norm(
-            encoder_decoder_attended,
-            self_attended,
-            self._enc_dec_attention_layernorm,
-        )
+
+        if self._cross_attention_layer is not None:
+            encoder_mask = merge_padding_and_attention_mask(
+                encoder_sequence, encoder_padding_mask, encoder_attention_mask
+            )
+            # Cross attention.
+            cross_attended = self._cross_attention_layer(
+                query=attention_output,
+                value=encoder_sequence,
+                key=encoder_sequence,
+                attention_mask=encoder_mask,
+            )
+            cross_attended = self._cross_attention_dropout(
+                cross_attended,
+            )
+            attention_output = self._add_and_norm(
+                cross_attended,
+                attention_output,
+                self._cross_attention_layernorm,
+            )
+
         # Feedforward.
-        feed_forward_output = self._feed_forward(encoder_decoder_attended)
+        feed_forward_output = self._feed_forward(attention_output)
         return self._add_and_norm(
-            encoder_decoder_attended,
+            attention_output,
             feed_forward_output,
             self._feedforward_layernorm,
         )


Test patch:
diff --git a/keras_nlp/layers/start_end_packer_test.py b/keras_nlp/layers/start_end_packer_test.py
new file mode 100644
index 00000000..b36bc3b1
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer_test.py
@@ -0,0 +1,151 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Start End Packer layer."""
+
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.start_end_packer import StartEndPacker
+
+
+class StartEndPackerTest(tf.test.TestCase):
+    def test_dense_input(self):
+        input_data = tf.constant([5, 6, 7])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [5, 6, 7, 0, 0]
+        self.assertAllEqual(output, expected_output)
+
+    def test_dense_input_error(self):
+        input_data = tf.constant([[5, 6, 7]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        with self.assertRaisesRegex(
+            ValueError,
+            "Input must either be dense with rank 1, or ragged with rank 2. "
+            "Received dense input with rank=2",
+        ):
+            start_end_packer(input_data)
+
+    def test_ragged_input(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [[5, 6, 7, 0, 0], [8, 9, 10, 11, 0]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_ragged_input_error(self):
+        input_data = tf.ragged.constant([[[5, 6, 7], [8, 9, 10, 11]]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        with self.assertRaisesRegex(
+            ValueError,
+            "Input must either be dense with rank 1, or ragged with rank 2. "
+            "Received ragged input with "
+            "rank=3",
+        ):
+            start_end_packer(input_data)
+
+    def test_start_end_token(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=6, start_value=1, end_value=2
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 0], [1, 8, 9, 10, 11, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_start_end_padding_value(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_end_token_value_during_truncation(self):
+        input_data = tf.ragged.constant([[5, 6], [8, 9, 10, 11, 12, 13]])
+        start_end_packer = StartEndPacker(
+            sequence_length=5, start_value=1, end_value=2, pad_value=0
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 2, 0], [1, 8, 9, 10, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_string_input(self):
+        input_data = tf.ragged.constant(
+            [["KerasNLP", "is", "awesome"], ["amazing"]]
+        )
+        start_end_packer = StartEndPacker(
+            sequence_length=5,
+            start_value="[START]",
+            end_value="[END]",
+            pad_value="[PAD]",
+        )
+        output = start_end_packer(input_data)
+        expected_output = [
+            ["[START]", "KerasNLP", "is", "awesome", "[END]"],
+            ["[START]", "amazing", "[END]", "[PAD]", "[PAD]"],
+        ]
+        self.assertAllEqual(output, expected_output)
+
+    def test_functional_model(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        inputs = keras.Input(dtype=tf.int32, shape=())
+        outputs = start_end_packer(inputs)
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(model_output, expected_output)
+
+    def test_batch(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        )
+        ds = ds.batch(2).map(start_end_packer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_get_config(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=512,
+            start_value=10,
+            end_value=20,
+            pad_value=100,
+            name="start_end_packer_test",
+        )
+
+        config = start_end_packer.get_config()
+        expected_config_subset = {
+            "sequence_length": 512,
+            "start_value": 10,
+            "end_value": 20,
+            "pad_value": 100,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 43460892..80a18e12 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -29,32 +29,52 @@ class TransformerDecoderTest(tf.test.TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        output = decoder(encoder_input, decoder_input)
+        output = decoder(decoder_input, encoder_input)
         model = keras.Model(
             inputs=[decoder_input, encoder_input],
             outputs=output,
         )
-
         encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        model([encoder_sequence, decoder_sequence])
+        model([decoder_sequence, encoder_sequence])
 
-    def test_valid_call_with_mask(self):
+    def test_valid_call_without_cross_attention(self):
+        decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
-        encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        output = decoder(decoder_input)
+        model = keras.Model(
+            inputs=decoder_input,
+            outputs=output,
+        )
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        encoder_padding_mask = encoder_sequence[:, :, 0] > 0.5
-        decoder_padding_mask = decoder_sequence[:, :, 0] > 0.5
-        decoder(
-            decoder_sequence,
-            encoder_sequence,
-            decoder_padding_mask=decoder_padding_mask,
-            encoder_padding_mask=encoder_padding_mask,
-            use_causal_mask=True,
+        model(decoder_sequence)
+
+    def test_invalid_calls(self):
+        encoder_input = keras.Input(shape=[4, 6])
+        decoder_input = keras.Input(shape=[4, 6])
+
+        # with cross-attention.
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder(decoder_input, encoder_input)
+        # should raise ValueError if encoder_input is not provided
+        with self.assertRaises(ValueError):
+            decoder(decoder_input)
+
+        # without cross-attention.
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
         )
+        decoder(decoder_input)
+        # should raise ValueError if encoder_input is provided
+        with self.assertRaises(ValueError):
+            decoder(decoder_input, encoder_input)
 
     def test_get_config_and_from_config(self):
         decoder = transformer_decoder.TransformerDecoder(
@@ -65,7 +85,6 @@ class TransformerDecoderTest(tf.test.TestCase):
         )
 
         config = decoder.get_config()
-
         expected_config_subset = {
             "intermediate_dim": 4,
             "num_heads": 2,
@@ -79,14 +98,11 @@ class TransformerDecoderTest(tf.test.TestCase):
                 keras.initializers.Zeros()
             ),
         }
-
         self.assertEqual(config, {**config, **expected_config_subset})
         self.assertEqual(config, {**config, **expected_config_subset})
-
         restored_decoder = transformer_decoder.TransformerDecoder.from_config(
             config,
         )
-
         self.assertEqual(
             restored_decoder.get_config(), {**config, **expected_config_subset}
         )
@@ -100,7 +116,7 @@ class TransformerDecoderTest(tf.test.TestCase):
                 kernel_initializer="Invalid",
             )
 
-    def test_one_training_step_of_transformer_encoder(self):
+    def test_one_training_step_of_transformer_with_cross_attention(self):
         class MyModel(keras.Model):
             def __init__(self):
                 super(MyModel, self).__init__()
@@ -128,12 +144,39 @@ class TransformerDecoderTest(tf.test.TestCase):
         self.assertGreater(len(grad), 1)
         optimizer.apply_gradients(zip(grad, model.trainable_variables))
 
+    def test_one_training_step_of_transformer_without_cross_attention(self):
+        class MyModel(keras.Model):
+            def __init__(self):
+                super(MyModel, self).__init__()
+                self._decoder = transformer_decoder.TransformerDecoder(
+                    intermediate_dim=4,
+                    num_heads=2,
+                )
+                self._dense = keras.layers.Dense(1, activation="sigmoid")
+
+            def call(self, decoder_input):
+                x = self._decoder(decoder_input)
+                return self._dense(x)
+
+        model = MyModel()
+
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        label = tf.cast(decoder_sequence[:, :, 0] >= 0.5, dtype=tf.int32)
+
+        loss_fn = keras.losses.BinaryCrossentropy(from_logits=False)
+        optimizer = keras.optimizers.Adam()
+        with tf.GradientTape() as tape:
+            pred = model(decoder_sequence)
+            loss = loss_fn(label, pred)
+        grad = tape.gradient(loss, model.trainable_variables)
+        self.assertGreater(len(grad), 1)
+        optimizer.apply_gradients(zip(grad, model.trainable_variables))
+
     def test_checkpointing_transformer_decoder(self):
         decoder1 = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
-
         decoder2 = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
@@ -156,10 +199,38 @@ class TransformerDecoderTest(tf.test.TestCase):
         checkpoint2.restore(save_path)
 
         decoder1_output = decoder1(decoder_sequence, encoder_sequence)
-        decoder2_output = decoder2(
-            decoder_sequence,
-            encoder_sequence,
+        decoder2_output = decoder2(decoder_sequence, encoder_sequence)
+        self.assertAllClose(decoder1_output, decoder2_output)
+
+    def test_checkpointing_transformer_decoder_without_cross_attention(self):
+        decoder1 = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
         )
+
+        decoder2 = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        decoder1(decoder_sequence)
+        decoder2(decoder_sequence)
+        # The weights of decoder1 and decoder2 are different.
+        self.assertFalse(
+            all(
+                decoder1._output_dense.trainable_variables[0][0]
+                == decoder2._output_dense.trainable_variables[0][0]
+            )
+        )
+        checkpoint = tf.train.Checkpoint(decoder1)
+        checkpoint2 = tf.train.Checkpoint(decoder2)
+        temp_dir = self.get_temp_dir()
+        save_path = checkpoint.save(temp_dir)
+        checkpoint2.restore(save_path)
+
+        decoder1_output = decoder1(decoder_sequence)
+        decoder2_output = decoder2(decoder_sequence)
         self.assertAllClose(decoder1_output, decoder2_output)
 
     def test_mask_propagation(self):
@@ -174,6 +245,17 @@ class TransformerDecoderTest(tf.test.TestCase):
         outputs = decoder(decoder_sequence, encoder_sequence)
         self.assertAllEqual(outputs._keras_mask, mask)
 
+    def test_mask_propagation_without_cross_attention(self):
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        decoder_sequence._keras_mask = mask
+        outputs = decoder(decoder_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])
@@ -188,11 +270,32 @@ class TransformerDecoderTest(tf.test.TestCase):
         )
         encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        model([encoder_sequence, decoder_sequence])
+        model([decoder_sequence, encoder_sequence])
         path = os.path.join(self.get_temp_dir(), "model")
         model.save(path)
-        loaded_model = keras.models.load_model(path)
 
+        loaded_model = keras.models.load_model(path)
         model_output = model([decoder_sequence, encoder_sequence])
         loaded_model_output = loaded_model([decoder_sequence, encoder_sequence])
         self.assertAllClose(model_output, loaded_model_output)
+
+    def test_save_model_without_cross_attention(self):
+        decoder_input = keras.Input(shape=[4, 6])
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        output = decoder(decoder_input)
+        model = keras.Model(
+            inputs=decoder_input,
+            outputs=output,
+        )
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        model(decoder_sequence)
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path)
+        loaded_model = keras.models.load_model(path)
+
+        model_output = model(decoder_sequence)
+        loaded_model_output = loaded_model(decoder_sequence)
+        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
index 80b86bc7..752bb4bd 100644
--- a/keras_nlp/tokenizers/byte_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -130,7 +130,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_tokenize_first_batch_second_with_sequence_length(self):
@@ -150,7 +149,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second(self):
@@ -169,7 +167,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second_with_sequence_length(self):
@@ -188,7 +185,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_functional_model(self):