Repository: keras-team/keras-nlp
Pull request #1240 (2023-09-09 01:19:10)

fixed token_to_id doc + error msg


Patch:
diff --git a/keras_nlp/layers/modeling/cached_multi_head_attention.py b/keras_nlp/layers/modeling/cached_multi_head_attention.py
index eb73ec4a..16124328 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention.py
@@ -138,4 +138,7 @@ class CachedMultiHeadAttention(keras.layers.MultiHeadAttention):
             self._combine_equation, attention_scores, value
         )
         attention_output = self._output_dense(attention_output)
-        return attention_output, cache
+
+        if cache is not None:
+            return attention_output, cache
+        return attention_output
diff --git a/keras_nlp/layers/modeling/f_net_encoder.py b/keras_nlp/layers/modeling/f_net_encoder.py
index aacd4d10..f8de353e 100644
--- a/keras_nlp/layers/modeling/f_net_encoder.py
+++ b/keras_nlp/layers/modeling/f_net_encoder.py
@@ -125,6 +125,7 @@ class FNetEncoder(keras.layers.Layer):
             self._intermediate_dense.compute_output_shape(inputs_shape)
         )
         self._output_dropout = keras.layers.Dropout(rate=self.dropout)
+        self.built = True
 
     def call(self, inputs):
         """Forward pass of the FNetEncoder.
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index 84186ab4..5bbcb491 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -167,8 +167,11 @@ class MaskedLMHead(keras.layers.Layer):
             initializer=self.bias_initializer,
             dtype=self.dtype,
         )
+        self.built = True
 
     def call(self, inputs, mask_positions):
+        # Avoid auto-converting numpy int arrays to float tensors.
+        mask_positions = ops.convert_to_tensor(mask_positions, dtype="int")
         # Gather the encoded tokens at the masked indices.
         mask_positions = ops.expand_dims(mask_positions, axis=-1)
         x = ops.take_along_axis(inputs, mask_positions, axis=1)
@@ -222,6 +225,4 @@ class MaskedLMHead(keras.layers.Layer):
         return config
 
     def compute_output_shape(self, inputs_shape, mask_positions_shape):
-        output_shape = list(mask_positions_shape)
-        output_shape[-1] = self.vocabulary_size
-        return tuple(output_shape)
+        return mask_positions_shape + (self.vocabulary_size,)
diff --git a/keras_nlp/layers/modeling/position_embedding.py b/keras_nlp/layers/modeling/position_embedding.py
index 6effdc6f..13a6e00a 100644
--- a/keras_nlp/layers/modeling/position_embedding.py
+++ b/keras_nlp/layers/modeling/position_embedding.py
@@ -92,16 +92,15 @@ class PositionEmbedding(keras.layers.Layer):
         )
         return config
 
-    def build(self, input_shape):
-        feature_size = input_shape[-1]
+    def build(self, inputs_shape):
+        feature_size = inputs_shape[-1]
         self.position_embeddings = self.add_weight(
             name="embeddings",
             shape=[self.sequence_length, feature_size],
             initializer=self.initializer,
             trainable=True,
         )
-
-        super().build(input_shape)
+        self.built = True
 
     def call(self, inputs, start_index=0):
         shape = ops.shape(inputs)
diff --git a/keras_nlp/layers/modeling/rotary_embedding.py b/keras_nlp/layers/modeling/rotary_embedding.py
index 5b5cb113..e5723b7c 100644
--- a/keras_nlp/layers/modeling/rotary_embedding.py
+++ b/keras_nlp/layers/modeling/rotary_embedding.py
@@ -82,6 +82,7 @@ class RotaryEmbedding(keras.layers.Layer):
         self.sequence_axis = sequence_axis
         self.feature_axis = feature_axis
         self.scaling_factor = scaling_factor
+        self.built = True
 
     def call(self, inputs, start_index=0):
         rotary_dim = ops.shape(inputs)[-1]
diff --git a/keras_nlp/layers/modeling/sine_position_encoding.py b/keras_nlp/layers/modeling/sine_position_encoding.py
index a5663e3b..84d042a7 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding.py
@@ -67,6 +67,7 @@ class SinePositionEncoding(keras.layers.Layer):
     ):
         super().__init__(**kwargs)
         self.max_wavelength = max_wavelength
+        self.built = True
 
     def call(self, inputs, start_index=0):
         shape = ops.shape(inputs)
diff --git a/keras_nlp/layers/modeling/transformer_decoder.py b/keras_nlp/layers/modeling/transformer_decoder.py
index 26c2e014..4117a82e 100644
--- a/keras_nlp/layers/modeling/transformer_decoder.py
+++ b/keras_nlp/layers/modeling/transformer_decoder.py
@@ -370,13 +370,17 @@ class TransformerDecoder(keras.layers.Layer):
         residual = x
         if self.normalize_first:
             x = self._self_attention_layer_norm(x)
-        x, self_attention_cache = self._self_attention_layer(
+        attention_output = self._self_attention_layer(
             query=x,
             value=x,
             attention_mask=self_attention_mask,
             cache=self_attention_cache,
             cache_update_index=self_attention_cache_update_index,
         )
+        if self_attention_cache is None:
+            x = attention_output
+        else:
+            x, self_attention_cache = attention_output
         x = self._self_attention_dropout(x)
         x = x + residual
         if not self.normalize_first:
@@ -393,13 +397,17 @@ class TransformerDecoder(keras.layers.Layer):
             residual = x
             if self.normalize_first:
                 x = self._cross_attention_layer_norm(x)
-            x, cross_attention_cache = self._cross_attention_layer(
+            attention_output = self._cross_attention_layer(
                 query=x,
                 value=encoder_sequence,
                 attention_mask=cross_attention_mask,
                 cache=cross_attention_cache,
                 cache_update_index=cross_attention_cache_update_index,
             )
+            if self_attention_cache is None:
+                x = attention_output
+            else:
+                x, self_attention_cache = attention_output
             x = self._cross_attention_dropout(x)
             x = x + residual
             if not self.normalize_first:
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
index e03123b0..bd6f568a 100644
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -115,9 +115,9 @@ class Tokenizer(PreprocessingLayer):
         )
 
     def token_to_id(self, token: str) -> int:
-        """Convert an integer id to a string token."""
+        """Convert a string token to an integer id."""
         raise NotImplementedError(
-            "No implementation of `id_to_token()` was found for "
+            "No implementation of `token_to_id()` was found for "
             f"{self.__class__.__name__}."
         )
 


Test patch:
diff --git a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
index 4d197504..3f6ecbc7 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
@@ -20,10 +20,21 @@ from keras_nlp.tests.test_case import TestCase
 
 
 class CachedMultiHeadAttentionTest(TestCase):
-    def test_valid_call(self):
-        layer = CachedMultiHeadAttention(num_heads=2, key_dim=4)
-        x = ops.random.uniform(shape=(2, 2, 8))
-        layer(query=x, value=x)
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=CachedMultiHeadAttention,
+            init_kwargs={
+                "num_heads": 2,
+                "key_dim": 4,
+            },
+            input_data={
+                "query": ops.random.uniform(shape=(2, 4, 6)),
+                "value": ops.random.uniform(shape=(2, 4, 6)),
+            },
+            expected_output_shape=(2, 4, 6),
+            expected_num_trainable_weights=8,
+            expected_num_non_trainable_variables=1,
+        )
 
     def test_cache_call_is_correct(self):
         batch_size = 2
diff --git a/keras_nlp/layers/modeling/f_net_encoder_test.py b/keras_nlp/layers/modeling/f_net_encoder_test.py
index d6254516..ffafc8a7 100644
--- a/keras_nlp/layers/modeling/f_net_encoder_test.py
+++ b/keras_nlp/layers/modeling/f_net_encoder_test.py
@@ -12,93 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import f_net_encoder
+from keras_nlp.layers.modeling.f_net_encoder import FNetEncoder
 from keras_nlp.tests.test_case import TestCase
 
 
 class FNetEncoderTest(TestCase):
-    def test_valid_call(self):
-        encoder = f_net_encoder.FNetEncoder(intermediate_dim=4)
-        model = keras.Sequential(
-            [
-                keras.Input(shape=(4, 6)),
-                encoder,
-            ]
-        )
-        input = ops.random.uniform(shape=[2, 4, 6])
-        model(input)
-
-    def test_get_config_and_from_config(self):
-        encoder = f_net_encoder.FNetEncoder(
-            intermediate_dim=4,
-            kernel_initializer="HeNormal",
-            bias_initializer="Zeros",
-        )
-        config = encoder.get_config()
-        expected_config_subset = {
-            "intermediate_dim": 4,
-            "dropout": 0,
-            "activation": "relu",
-            "layer_norm_epsilon": 1e-5,
-            "kernel_initializer": keras.initializers.serialize(
-                keras.initializers.HeNormal()
-            ),
-            "bias_initializer": keras.initializers.serialize(
-                keras.initializers.Zeros()
-            ),
-        }
-        self.assertEqual(config, {**config, **expected_config_subset})
-
-        restored_encoder = f_net_encoder.FNetEncoder.from_config(
-            config,
-        )
-        self.assertEqual(
-            restored_encoder.get_config(), {**config, **expected_config_subset}
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=FNetEncoder,
+            init_kwargs={
+                "intermediate_dim": 4,
+                "dropout": 0,
+                "activation": "relu",
+                "layer_norm_epsilon": 1e-5,
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+            },
+            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            expected_output_shape=(2, 4, 6),
+            expected_num_trainable_weights=8,
+            expected_num_non_trainable_variables=1,
         )
 
     def test_value_error_when_invalid_kernel_initializer(self):
         with self.assertRaises(ValueError):
-            f_net_encoder.FNetEncoder(
+            FNetEncoder(
                 intermediate_dim=4,
                 dropout=0.5,
                 kernel_initializer="Invalid",
             )
-
-    def test_one_training_step_of_f_net_encoder(self):
-        encoder = f_net_encoder.FNetEncoder(intermediate_dim=4)
-        inputs = keras.Input(shape=(4, 6))
-        x = encoder(inputs)
-        x = keras.layers.Dense(1, activation="sigmoid")(x)
-        model = keras.Model(inputs=inputs, outputs=x)
-
-        data = ops.random.uniform(shape=[2, 4, 6])
-        label = ops.random.randint(minval=0, maxval=2, shape=(2, 4, 1))
-
-        loss = keras.losses.BinaryCrossentropy(from_logits=False)
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(x=data, y=label)
-        self.assertGreater(loss, 0)
-
-    def test_saved_model(self):
-        model = keras.Sequential(
-            [
-                keras.Input(shape=(4, 6)),
-                f_net_encoder.FNetEncoder(
-                    intermediate_dim=4,
-                ),
-            ]
-        )
-        data = ops.random.uniform(shape=[2, 4, 6])
-        model(data)
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        loaded_model = keras.models.load_model(path)
-
-        model_output = model(data)
-        loaded_model_output = loaded_model(data)
-        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/layers/modeling/masked_lm_head_test.py b/keras_nlp/layers/modeling/masked_lm_head_test.py
index 7f08f292..9ccdcae1 100644
--- a/keras_nlp/layers/modeling/masked_lm_head_test.py
+++ b/keras_nlp/layers/modeling/masked_lm_head_test.py
@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.masked_lm_head import MaskedLMHead
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
@@ -22,70 +19,46 @@ from keras_nlp.tests.test_case import TestCase
 
 
 class MaskedLMHeadTest(TestCase):
-    def test_valid_call(self):
-        head = MaskedLMHead(
-            vocabulary_size=100,
-            activation="softmax",
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=MaskedLMHead,
+            init_kwargs={
+                "vocabulary_size": 100,
+                "activation": "softmax",
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+            },
+            input_data={
+                "inputs": ops.random.uniform(shape=(4, 10, 16)),
+                "mask_positions": ops.random.randint(
+                    minval=0, maxval=10, shape=(4, 5)
+                ),
+            },
+            expected_output_shape=(4, 5, 100),
+            expected_num_trainable_weights=6,
         )
-        encoded_tokens = keras.Input(shape=(10, 16))
-        positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, mask_positions=positions)
-        model = keras.Model((encoded_tokens, positions), outputs)
-
-        token_data = ops.random.uniform(shape=(4, 10, 16))
-        position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
-        model((token_data, position_data))
 
-    def test_valid_call_with_token_embedding(self):
+    def test_layer_behaviors_with_embedding(self):
         embedding = ReversibleEmbedding(100, 16)
         embedding.build((4, 10))
-        head = MaskedLMHead(
-            vocabulary_size=100,
-            token_embedding=embedding,
-            activation="softmax",
-        )
-        # Use a difference "hidden dim" for the model than "embedding dim", we
-        # need to support this in the layer.
-        sequence = keras.Input(shape=(10, 32))
-        positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(sequence, mask_positions=positions)
-        model = keras.Model((sequence, positions), outputs)
-        sequence_data = ops.random.uniform(shape=(4, 10, 32))
-        position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
-        model((sequence_data, position_data))
-
-    def test_get_config_and_from_config(self):
-        head = MaskedLMHead(
-            vocabulary_size=100,
-            kernel_initializer="HeNormal",
-            bias_initializer="Zeros",
-            activation="softmax",
-        )
-
-        config = head.get_config()
-
-        expected_params = {
-            "vocabulary_size": 100,
-            "kernel_initializer": keras.initializers.serialize(
-                keras.initializers.HeNormal()
-            ),
-            "bias_initializer": keras.initializers.serialize(
-                keras.initializers.Zeros()
-            ),
-            "activation": keras.activations.serialize(
-                keras.activations.softmax
-            ),
-        }
-
-        self.assertEqual(config, {**config, **expected_params})
-
-        restored = MaskedLMHead.from_config(config)
-        restored_config = restored.get_config()
-
-        self.assertEqual(
-            restored_config, {**restored_config, **expected_params}
+        self.run_layer_test(
+            layer_cls=MaskedLMHead,
+            init_kwargs={
+                "vocabulary_size": 100,
+                "activation": "softmax",
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+                "token_embedding": embedding,
+            },
+            input_data={
+                "inputs": ops.random.uniform(shape=(4, 10, 16)),
+                "mask_positions": ops.random.randint(
+                    minval=0, maxval=10, shape=(4, 5)
+                ),
+            },
+            expected_output_shape=(4, 5, 100),
+            expected_num_trainable_weights=6,
         )
-        self.assertEqual(restored_config, config)
 
     def test_value_error_when_neither_embedding_or_vocab_size_set(self):
         with self.assertRaises(ValueError):
@@ -99,42 +72,3 @@ class MaskedLMHeadTest(TestCase):
                 vocabulary_size=101,
                 token_embedding=embedding,
             )
-
-    def test_one_train_step(self):
-        head = MaskedLMHead(
-            vocabulary_size=100,
-        )
-        encoded_tokens = keras.Input(shape=(10, 16))
-        positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, mask_positions=positions)
-        model = keras.Model((encoded_tokens, positions), outputs)
-
-        token_data = ops.random.uniform(shape=(4, 10, 16))
-        position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
-        label_data = ops.random.randint(minval=0, maxval=2, shape=(4, 5, 1))
-
-        loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(x=(token_data, position_data), y=label_data)
-        self.assertGreater(loss, 0)
-
-    def test_saved_model(self):
-        head = MaskedLMHead(
-            vocabulary_size=100,
-            activation="softmax",
-        )
-        encoded_tokens = keras.Input(shape=(10, 16))
-        positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, mask_positions=positions)
-        model = keras.Model((encoded_tokens, positions), outputs)
-
-        token_data = ops.random.uniform(shape=(4, 10, 16))
-        position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
-        model_output = model((token_data, position_data))
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        restored_model = keras.models.load_model(path)
-
-        restored_output = restored_model((token_data, position_data))
-        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/layers/modeling/position_embedding_test.py b/keras_nlp/layers/modeling/position_embedding_test.py
index f3010d80..ec16f541 100644
--- a/keras_nlp/layers/modeling/position_embedding_test.py
+++ b/keras_nlp/layers/modeling/position_embedding_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 import numpy as np
 
 from keras_nlp.backend import keras
@@ -30,42 +28,27 @@ def custom_init(shape, dtype=None):
 
 
 class PositionEmbeddingTest(TestCase):
-    def test_static_layer_output_shape(self):
-        # Create a 3-dimensional input (the first dimension is implicit).
-        sequence_length = 21
-        feature_size = 30
-        test_layer = PositionEmbedding(sequence_length=sequence_length)
-        input_tensor = keras.Input(shape=(sequence_length, feature_size))
-        output_tensor = test_layer(input_tensor)
-
-        # When using static position embedding shapes, the output is expected
-        # to be the same as the input shape in all dimensions save batch.
-        expected_output_shape = (None, sequence_length, feature_size)
-        self.assertEqual(expected_output_shape, output_tensor.shape)
-        # The output dtype for this layer should match the compute dtype.
-        self.assertEqual(test_layer.compute_dtype, output_tensor.dtype)
-
-    def test_more_than_3_dimensions_static(self):
-        # Create a 4-dimensional input (the first dimension is implicit).
-        sequence_length = 21
-        feature_size = 30
-        test_layer = PositionEmbedding(sequence_length=sequence_length)
-        input_tensor = keras.Input(
-            shape=(feature_size, sequence_length, feature_size)
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=PositionEmbedding,
+            init_kwargs={
+                "sequence_length": 21,
+            },
+            input_data=ops.random.uniform(shape=(4, 21, 30)),
+            expected_output_shape=(4, 21, 30),
+            expected_num_trainable_weights=1,
         )
-        output_tensor = test_layer(input_tensor)
 
-        # When using static position embedding shapes, the output is expected
-        # to be the same as the input shape in all dimensions save batch.
-        expected_output_shape = (
-            None,
-            feature_size,
-            sequence_length,
-            feature_size,
+    def test_layer_behaviors_4d(self):
+        self.run_layer_test(
+            layer_cls=PositionEmbedding,
+            init_kwargs={
+                "sequence_length": 21,
+            },
+            input_data=ops.random.uniform(shape=(4, 5, 21, 30)),
+            expected_output_shape=(4, 5, 21, 30),
+            expected_num_trainable_weights=1,
         )
-        self.assertEqual(expected_output_shape, output_tensor.shape)
-        # The output dtype for this layer should match the compute dtype.
-        self.assertEqual(test_layer.compute_dtype, output_tensor.dtype)
 
     def test_float16_dtype(self):
         # Create a 3-dimensional input (the first dimension is implicit).
@@ -171,54 +154,3 @@ class PositionEmbeddingTest(TestCase):
                 sequential_output, (0, i, 0), parial_output
             )
         self.assertAllClose(full_output, sequential_output)
-
-    def test_one_training_step(self):
-        max_sequence_length = 4
-        feature_size = 3
-        inputs = keras.Input(shape=(max_sequence_length, feature_size))
-        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
-        outputs = test_layer(inputs)
-        model = keras.Model(inputs=inputs, outputs=outputs)
-
-        batch_size = 2
-        data = ops.random.uniform(
-            shape=[batch_size, max_sequence_length, feature_size]
-        )
-        label = ops.random.uniform(
-            shape=[batch_size, max_sequence_length, feature_size]
-        )
-
-        loss = keras.losses.MeanSquaredError()
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(x=data, y=label)
-        self.assertGreater(loss, 0)
-
-    def test_get_config_and_from_config(self):
-        max_sequence_length = 40
-        test_layer = PositionEmbedding(
-            sequence_length=max_sequence_length,
-            initializer="zeros",
-        )
-        config = test_layer.get_config()
-        restored = PositionEmbedding.from_config(config)
-        self.assertEqual(restored.get_config(), config)
-
-    def test_saved_model(self):
-        max_sequence_length = 4
-        feature_size = 6
-        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
-        inputs = keras.Input(shape=(max_sequence_length, feature_size))
-        outputs = test_layer(inputs)
-        model = keras.Model(inputs=inputs, outputs=outputs)
-
-        data = np.zeros(shape=[2, max_sequence_length, feature_size])
-        model(data)
-
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        loaded_model = keras.models.load_model(path)
-
-        model_output = model.predict(data)
-        loaded_model_output = loaded_model.predict(data)
-        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/layers/modeling/reversible_embedding_test.py b/keras_nlp/layers/modeling/reversible_embedding_test.py
index 4dd7026a..16cc3f8f 100644
--- a/keras_nlp/layers/modeling/reversible_embedding_test.py
+++ b/keras_nlp/layers/modeling/reversible_embedding_test.py
@@ -28,15 +28,19 @@ class ReversibleEmbeddingTest(TestCase):
         ("tie_weights", True),
         ("untie_weights", False),
     )
-    def test_valid_call(self, tie_weights):
-        embedding = ReversibleEmbedding(100, 32, tie_weights=tie_weights)
-        inputs = keras.Input(shape=(10,))
-        hidden_states = embedding(inputs)
-        outputs = embedding(hidden_states, reverse=True)
-        model = keras.Model(inputs, outputs)
-
-        input_data = ops.random.uniform(shape=(4, 10))
-        model(input_data)
+    def test_layer_behaviors_tied(self, tie_weights):
+        self.run_layer_test(
+            layer_cls=ReversibleEmbedding,
+            init_kwargs={
+                "input_dim": 100,
+                "output_dim": 32,
+                "tie_weights": tie_weights,
+                "embeddings_initializer": "HeNormal",
+            },
+            input_data=ops.random.randint(minval=0, maxval=100, shape=(4, 10)),
+            expected_output_shape=(4, 10, 32),
+            expected_num_trainable_weights=1 if tie_weights else 2,
+        )
 
     def test_correctness(self):
         layer = ReversibleEmbedding(input_dim=3, output_dim=2)
@@ -51,17 +55,6 @@ class ReversibleEmbeddingTest(TestCase):
         out = layer(np.array(([[1.0, 1.0]])), reverse=True)
         self.assertAllClose(out, np.array([[0.0, 4.0, 6.0]]))
 
-    def test_config(self):
-        original = ReversibleEmbedding(
-            100,
-            32,
-            tie_weights=False,
-            embeddings_initializer="HeNormal",
-        )
-        restored = ReversibleEmbedding.from_config(original.get_config())
-        restored.set_weights(original.get_weights())
-        self.assertEqual(restored.get_config(), original.get_config())
-
     def test_tied_checkpoint_untied_weights(self):
         embedding = ReversibleEmbedding(100, 16, tie_weights=True)
         inputs = keras.Input(shape=(10,), dtype="int32")
@@ -80,23 +73,3 @@ class ReversibleEmbeddingTest(TestCase):
 
         input_data = ops.ones(shape=(4, 10), dtype="int32")
         self.assertAllClose(untied_model(input_data), tied_model(input_data))
-
-    @parameterized.named_parameters(
-        ("tie_weights", True),
-        ("untie_weights", False),
-    )
-    def test_saved_model(self, tie_weights):
-        embedding = ReversibleEmbedding(100, 16, tie_weights=tie_weights)
-        inputs = keras.Input(shape=(10,), dtype="int32")
-        hidden_states = embedding(inputs)
-        outputs = embedding(hidden_states, reverse=True)
-        model = keras.Model(inputs, outputs)
-
-        input_data = ops.ones(shape=(4, 10), dtype="int32")
-        model_output = model(input_data)
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        restored_model = keras.models.load_model(path)
-
-        restored_output = restored_model(input_data)
-        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/layers/modeling/rotary_embedding_test.py b/keras_nlp/layers/modeling/rotary_embedding_test.py
index f0f0e25d..3fdac028 100644
--- a/keras_nlp/layers/modeling/rotary_embedding_test.py
+++ b/keras_nlp/layers/modeling/rotary_embedding_test.py
@@ -19,28 +19,28 @@ from keras_nlp.tests.test_case import TestCase
 
 
 class RotaryEmbeddingTest(TestCase):
-    def test_valid_call(self):
-        embedding_layer = RotaryEmbedding()
-        model = keras.Sequential(
-            [
-                keras.Input(shape=(4, 6)),
-                embedding_layer,
-            ]
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=RotaryEmbedding,
+            init_kwargs={
+                "max_wavelength": 1000,
+                "scaling_factor": 2.0,
+                "sequence_axis": 1,
+                "feature_axis": -1,
+            },
+            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            expected_output_shape=(2, 4, 6),
         )
-        input = ops.ones(shape=[2, 4, 6])
-        model(input)
 
-    def test_static_layer_output_shape(self):
-        embedding_layer = RotaryEmbedding()
-        seq_length = 100
-        hidden_size = 32
-        inputs = keras.Input(shape=(seq_length, hidden_size))
-        outputs = embedding_layer(inputs)
-
-        # When using static positional encoding shapes, the output is expected
-        # to be the same as the input shape in all dimensions.
-        expected_output_shape = (None, seq_length, hidden_size)
-        self.assertEqual(expected_output_shape, outputs.shape)
+    def test_layer_behaviors_4d(self):
+        self.run_layer_test(
+            layer_cls=RotaryEmbedding,
+            init_kwargs={
+                "max_wavelength": 1000,
+            },
+            input_data=ops.random.uniform(shape=(2, 8, 4, 6)),
+            expected_output_shape=(2, 8, 4, 6),
+        )
 
     def test_dynamic_layer_output_shape(self):
         embedding_layer = RotaryEmbedding()
@@ -96,27 +96,6 @@ class RotaryEmbeddingTest(TestCase):
             )
         self.assertAllClose(full_output, sequential_output)
 
-    def test_get_config_and_from_config(self):
-        embedding_layer = RotaryEmbedding(
-            max_wavelength=1000,
-            scaling_factor=2.0,
-            sequence_axis=1,
-            feature_axis=-1,
-        )
-        config = embedding_layer.get_config()
-        expected_config = {
-            "max_wavelength": 1000,
-            "scaling_factor": 2.0,
-            "sequence_axis": 1,
-            "feature_axis": -1,
-        }
-        self.assertEqual(config, {**config, **expected_config})
-        restored_embedding_layer = RotaryEmbedding.from_config(config)
-        self.assertEqual(
-            restored_embedding_layer.get_config(),
-            {**config, **expected_config},
-        )
-
     def test_float16_dtype(self):
         embedding_layer = RotaryEmbedding(dtype="float16")
         seq_length = 100
diff --git a/keras_nlp/layers/modeling/sine_position_encoding_test.py b/keras_nlp/layers/modeling/sine_position_encoding_test.py
index fc04d81e..22d1d9d3 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding_test.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding_test.py
@@ -21,16 +21,25 @@ from keras_nlp.tests.test_case import TestCase
 
 
 class SinePositionEncodingTest(TestCase):
-    def test_valid_call(self):
-        pos_encoding = SinePositionEncoding()
-        model = keras.Sequential(
-            [
-                keras.Input(shape=(4, 6)),
-                pos_encoding,
-            ]
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=SinePositionEncoding,
+            init_kwargs={
+                "max_wavelength": 10000,
+            },
+            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            expected_output_shape=(2, 4, 6),
+        )
+
+    def test_layer_behaviors_4d(self):
+        self.run_layer_test(
+            layer_cls=SinePositionEncoding,
+            init_kwargs={
+                "max_wavelength": 10000,
+            },
+            input_data=ops.random.uniform(shape=(1, 2, 4, 6)),
+            expected_output_shape=(1, 2, 4, 6),
         )
-        input = ops.random.uniform(shape=[2, 4, 6])
-        model(input)
 
     def test_static_layer_output_shape(self):
         pos_encoding = SinePositionEncoding()
@@ -98,21 +107,6 @@ class SinePositionEncodingTest(TestCase):
             )
         self.assertAllClose(full_output, sequential_output)
 
-    def test_get_config_and_from_config(self):
-        pos_encoding = SinePositionEncoding(
-            max_wavelength=1000,
-        )
-        config = pos_encoding.get_config()
-        expected_config_subset = {
-            "max_wavelength": 1000,
-        }
-        self.assertEqual(config, {**config, **expected_config_subset})
-        restored_pos_encoding = SinePositionEncoding.from_config(config)
-        self.assertEqual(
-            restored_pos_encoding.get_config(),
-            {**config, **expected_config_subset},
-        )
-
     def test_float16_dtype(self):
         pos_encoding = SinePositionEncoding(dtype="float16")
         seq_length = 100
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding_test.py b/keras_nlp/layers/modeling/token_and_position_embedding_test.py
index 415d845e..ceb96b11 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding_test.py
@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 import numpy as np
 
 from keras_nlp.backend import keras
+from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
@@ -24,54 +23,21 @@ from keras_nlp.tests.test_case import TestCase
 
 
 class TokenAndPositionEmbeddingTest(TestCase):
-    def test_get_config_and_from_config(self):
-        token_and_position_embed = TokenAndPositionEmbedding(
-            vocabulary_size=5,
-            sequence_length=10,
-            embedding_dim=32,
-        )
-
-        config = token_and_position_embed.get_config()
-
-        expected_config_subset = {
-            "embeddings_initializer": keras.initializers.serialize(
-                keras.initializers.GlorotUniform()
-            ),
-            "mask_zero": False,
-        }
-
-        self.assertEqual(config, {**config, **expected_config_subset})
-
-        restored_token_and_position_embed = (
-            TokenAndPositionEmbedding.from_config(config)
+    def test_layer_behaviors(self):
+        self.run_layer_test(
+            layer_cls=TokenAndPositionEmbedding,
+            init_kwargs={
+                "vocabulary_size": 5,
+                "sequence_length": 4,
+                "embedding_dim": 3,
+                "embeddings_initializer": keras.initializers.Constant(1.0),
+            },
+            input_data=ops.random.randint(minval=0, maxval=5, shape=(2, 4)),
+            expected_output_shape=(2, 4, 3),
+            expected_output_data=ops.ones((2, 4, 3)) * 2,
+            expected_num_trainable_weights=2,
         )
 
-        self.assertEqual(
-            restored_token_and_position_embed.get_config(),
-            {**config, **expected_config_subset},
-        )
-
-    def test_dense_tensor(self):
-        vocabulary_size = 5
-        sequence_length = 4
-        embedding_dim = 3
-        test_layer = TokenAndPositionEmbedding(
-            vocabulary_size=vocabulary_size,
-            sequence_length=sequence_length,
-            embedding_dim=embedding_dim,
-            embeddings_initializer=keras.initializers.Constant(1.0),
-        )
-        # Create a 2-dimensional input
-        # (the first dimension is implicit).
-        inputs = keras.Input(shape=(sequence_length,), dtype="int32")
-        outputs = test_layer(inputs)
-        model = keras.Model(inputs, outputs)
-
-        input_data = np.ones((2, sequence_length), dtype="int32")
-        expected_output_data = np.ones((2, sequence_length, embedding_dim)) * 2
-        output_data = model.predict(input_data)
-        self.assertAllClose(output_data, expected_output_data)
-
     def test_mask_propagation(self):
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=5,
@@ -83,27 +49,3 @@ class TokenAndPositionEmbeddingTest(TestCase):
         mask = input_data != 0
         outputs = test_layer(input_data)
         self.assertAllEqual(outputs._keras_mask, mask)
-
-    def test_saved_model(self):
-        vocabulary_size = 5
-        sequence_length = 4
-        embedding_dim = 3
-        test_layer = TokenAndPositionEmbedding(
-            vocabulary_size=vocabulary_size,
-            sequence_length=sequence_length,
-            embedding_dim=embedding_dim,
-        )
-        inputs = keras.Input(shape=(sequence_length,))
-        outputs = test_layer(inputs)
-        model = keras.Model(inputs=inputs, outputs=outputs)
-
-        data = np.zeros(shape=[2, sequence_length])
-        model(data)
-
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        loaded_model = keras.models.load_model(path)
-
-        model_output = model.predict(data)
-        loaded_model_output = loaded_model.predict(data)
-        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/layers/modeling/transformer_decoder_test.py b/keras_nlp/layers/modeling/transformer_decoder_test.py
index e24fd1ab..12cd189e 100644
--- a/keras_nlp/layers/modeling/transformer_decoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_decoder_test.py
@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 from absl.testing import parameterized
 
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import transformer_decoder
+from keras_nlp.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_nlp.tests.test_case import TestCase
 
 
@@ -27,48 +24,56 @@ class TransformerDecoderTest(TestCase):
         ("without_norm_first", False),
         ("with_norm_first", True),
     )
-    def test_valid_call(self, normalize_first):
-        encoder_input = keras.Input(shape=[4, 6])
-        decoder_input = keras.Input(shape=[4, 6])
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-            normalize_first=normalize_first,
+    def test_layer_behaviors(self, normalize_first):
+        self.run_layer_test(
+            layer_cls=TransformerDecoder,
+            init_kwargs={
+                "intermediate_dim": 4,
+                "num_heads": 2,
+                "normalize_first": normalize_first,
+                "activation": "relu",
+                "layer_norm_epsilon": 1e-05,
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+            },
+            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            expected_output_shape=(2, 4, 6),
+            expected_num_trainable_weights=16,
+            expected_num_non_trainable_variables=3,  # dropout rng seeds
         )
-        output = decoder(decoder_input, encoder_input)
-        model = keras.Model(
-            inputs=[decoder_input, encoder_input],
-            outputs=output,
-        )
-        encoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        decoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        model([decoder_sequence, encoder_sequence])
 
     @parameterized.named_parameters(
         ("without_norm_first", False),
         ("with_norm_first", True),
     )
-    def test_valid_call_without_cross_attention(self, normalize_first):
-        decoder_input = keras.Input(shape=[4, 6])
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-            normalize_first=normalize_first,
-        )
-        output = decoder(decoder_input)
-        model = keras.Model(
-            inputs=decoder_input,
-            outputs=output,
+    def test_layer_behaviors_with_cross_attention(self, normalize_first):
+        pass
+        self.run_layer_test(
+            layer_cls=TransformerDecoder,
+            init_kwargs={
+                "intermediate_dim": 4,
+                "num_heads": 2,
+                "normalize_first": normalize_first,
+                "activation": "relu",
+                "layer_norm_epsilon": 1e-05,
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+            },
+            input_data={
+                "decoder_sequence": ops.random.uniform(shape=(2, 4, 6)),
+                "encoder_sequence": ops.random.uniform(shape=(2, 4, 6)),
+            },
+            expected_output_shape=(2, 4, 6),
+            expected_num_trainable_weights=26,
+            expected_num_non_trainable_variables=5,  # dropout rng seeds
         )
-        decoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        model(decoder_sequence)
 
     def test_invalid_calls(self):
         encoder_input = ops.zeros((2, 4, 6))
         decoder_input = ops.zeros((2, 4, 6))
 
         # with cross-attention.
-        decoder = transformer_decoder.TransformerDecoder(
+        decoder = TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -78,7 +83,7 @@ class TransformerDecoderTest(TestCase):
             decoder(decoder_input)
 
         # without cross-attention.
-        decoder = transformer_decoder.TransformerDecoder(
+        decoder = TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -87,90 +92,17 @@ class TransformerDecoderTest(TestCase):
         with self.assertRaises(ValueError):
             decoder(decoder_input, encoder_input)
 
-    def test_get_config_and_from_config(self):
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-            kernel_initializer="HeNormal",
-            bias_initializer="Zeros",
-            normalize_first=True,
-        )
-
-        config = decoder.get_config()
-        expected_config_subset = {
-            "intermediate_dim": 4,
-            "num_heads": 2,
-            "dropout": 0,
-            "activation": "relu",
-            "layer_norm_epsilon": 1e-05,
-            "kernel_initializer": keras.initializers.serialize(
-                keras.initializers.HeNormal()
-            ),
-            "bias_initializer": keras.initializers.serialize(
-                keras.initializers.Zeros()
-            ),
-            "normalize_first": True,
-        }
-        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}
-        )
-
     def test_value_error_when_invalid_kernel_inititalizer(self):
         with self.assertRaises(ValueError):
-            transformer_decoder.TransformerDecoder(
+            TransformerDecoder(
                 intermediate_dim=4,
                 num_heads=2,
                 dropout=0.5,
                 kernel_initializer="Invalid",
             )
 
-    def test_one_training_step_of_transformer_with_cross_attention(self):
-        decoder_input = keras.Input(shape=(4, 6))
-        encoder_input = keras.Input(shape=(4, 6))
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4, num_heads=2
-        )
-        outputs = decoder(decoder_input, encoder_input)
-        outputs = keras.layers.Dense(10, activation="softmax")(outputs)
-        model = keras.Model((decoder_input, encoder_input), outputs)
-
-        decoder_sequence = ops.random.uniform(shape=(2, 4, 6))
-        encoder_sequence = ops.random.uniform(shape=(2, 4, 6))
-        label = ops.random.randint(minval=0, maxval=10, shape=(2, 4, 1))
-
-        loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(
-            x=(decoder_sequence, encoder_sequence), y=label
-        )
-        self.assertGreater(loss, 0)
-
-    def test_one_training_step_of_transformer_without_cross_attention(self):
-        decoder_input = keras.Input(shape=(4, 6))
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4, num_heads=2
-        )
-        outputs = decoder(decoder_input)
-        outputs = keras.layers.Dense(10, activation="softmax")(outputs)
-        model = keras.Model(decoder_input, outputs)
-
-        decoder_sequence = ops.random.uniform(shape=(2, 4, 6))
-        label = ops.random.randint(minval=0, maxval=10, shape=(2, 4, 1))
-
-        loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(x=decoder_sequence, y=label)
-        self.assertGreater(loss, 0)
-
     def test_mask_propagation(self):
-        decoder = transformer_decoder.TransformerDecoder(
+        decoder = TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -182,7 +114,7 @@ class TransformerDecoderTest(TestCase):
         self.assertAllEqual(outputs._keras_mask, mask)
 
     def test_mask_propagation_without_cross_attention(self):
-        decoder = transformer_decoder.TransformerDecoder(
+        decoder = TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -204,7 +136,7 @@ class TransformerDecoderTest(TestCase):
         input_cache = ops.zeros((batch_size, 2, seq_len, num_heads, key_dim))
         outputs = ops.zeros_like(x)
 
-        layer = transformer_decoder.TransformerDecoder(
+        layer = TransformerDecoder(
             intermediate_dim=4,
             num_heads=num_heads,
         )
@@ -234,52 +166,5 @@ class TransformerDecoderTest(TestCase):
             return outputs, cache
 
         output, output_cache = call(outputs, input_cache)
-
         self.assertAllClose(output, no_loop_outputs)
         self.assertAllClose(output_cache, no_loop_cache)
-
-    def test_saved_model(self):
-        encoder_input = keras.Input(shape=[4, 6])
-        decoder_input = keras.Input(shape=[4, 6])
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-            normalize_first=True,
-        )
-        output = decoder(encoder_input, decoder_input)
-        model = keras.Model(
-            inputs=[decoder_input, encoder_input],
-            outputs=output,
-        )
-        encoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        decoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        model([decoder_sequence, encoder_sequence])
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-
-        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_saved_model_without_cross_attention(self):
-        decoder_input = keras.Input(shape=[4, 6])
-        decoder = transformer_decoder.TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-            normalize_first=True,
-        )
-        output = decoder(decoder_input)
-        model = keras.Model(
-            inputs=decoder_input,
-            outputs=output,
-        )
-        decoder_sequence = ops.random.uniform(shape=[2, 4, 6])
-        model(decoder_sequence)
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-        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/layers/modeling/transformer_encoder_test.py b/keras_nlp/layers/modeling/transformer_encoder_test.py
index b8ea4523..9fe7c3ea 100644
--- a/keras_nlp/layers/modeling/transformer_encoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_encoder_test.py
@@ -12,23 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 from absl.testing import parameterized
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import transformer_encoder
+from keras_nlp.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_nlp.tests.test_case import TestCase
 
 
 class TransformerEncoderTest(TestCase):
+    @parameterized.named_parameters(
+        ("without_norm_first", False),
+        ("with_norm_first", True),
+    )
+    def test_layer_behaviors(self, normalize_first):
+        self.run_layer_test(
+            layer_cls=TransformerEncoder,
+            init_kwargs={
+                "intermediate_dim": 4,
+                "num_heads": 2,
+                "normalize_first": normalize_first,
+                "activation": "relu",
+                "layer_norm_epsilon": 1e-05,
+                "kernel_initializer": "HeNormal",
+                "bias_initializer": "Zeros",
+            },
+            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            expected_output_shape=(2, 4, 6),
+            expected_num_trainable_weights=16,
+            expected_num_non_trainable_variables=3,  # dropout rng seeds
+        )
+
     @parameterized.named_parameters(
         ("without_norm_first", False),
         ("with_norm_first", True),
     )
     def test_valid_call(self, normalize_first):
-        encoder = transformer_encoder.TransformerEncoder(
+        encoder = TransformerEncoder(
             intermediate_dim=4,
             num_heads=2,
             normalize_first=normalize_first,
@@ -43,7 +63,7 @@ class TransformerEncoderTest(TestCase):
         model(input)
 
     def test_valid_call_with_mask(self):
-        encoder = transformer_encoder.TransformerEncoder(
+        encoder = TransformerEncoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -52,72 +72,17 @@ class TransformerEncoderTest(TestCase):
         mask = input[:, :, 0] < 0.5
         encoder(input, mask)
 
-    def test_get_config_and_from_config(self):
-        encoder = transformer_encoder.TransformerEncoder(
-            intermediate_dim=4,
-            num_heads=2,
-            kernel_initializer="HeNormal",
-            bias_initializer="Zeros",
-            normalize_first=True,
-        )
-
-        config = encoder.get_config()
-
-        expected_config_subset = {
-            "intermediate_dim": 4,
-            "num_heads": 2,
-            "dropout": 0,
-            "activation": "relu",
-            "layer_norm_epsilon": 1e-05,
-            "kernel_initializer": keras.initializers.serialize(
-                keras.initializers.HeNormal()
-            ),
-            "bias_initializer": keras.initializers.serialize(
-                keras.initializers.Zeros()
-            ),
-            "normalize_first": True,
-        }
-
-        self.assertEqual(config, {**config, **expected_config_subset})
-
-        restored_encoder = transformer_encoder.TransformerEncoder.from_config(
-            config,
-        )
-
-        self.assertEqual(
-            restored_encoder.get_config(), {**config, **expected_config_subset}
-        )
-
     def test_value_error_when_invalid_kernel_inititalizer(self):
         with self.assertRaises(ValueError):
-            transformer_encoder.TransformerEncoder(
+            TransformerEncoder(
                 intermediate_dim=4,
                 num_heads=2,
                 dropout=0.5,
                 kernel_initializer="Invalid",
             )
 
-    def test_one_training_step_of_transformer_encoder(self):
-        encoder = transformer_encoder.TransformerEncoder(
-            intermediate_dim=4,
-            num_heads=2,
-        )
-        inputs = keras.Input(shape=(4, 6))
-        x = encoder(inputs)
-        x = keras.layers.Dense(1, activation="sigmoid")(x)
-        model = keras.Model(inputs=inputs, outputs=x)
-
-        data = ops.random.uniform(shape=[2, 4, 6])
-        label = ops.random.uniform(minval=0, maxval=2, shape=[2, 4, 1])
-
-        loss = keras.losses.MeanSquaredError()
-        optimizer = keras.optimizers.Adam()
-        model.compile(loss=loss, optimizer=optimizer)
-        loss = model.train_on_batch(x=data, y=label)
-        self.assertGreater(loss, 0)
-
     def test_mask_propagation(self):
-        encoder = transformer_encoder.TransformerEncoder(
+        encoder = TransformerEncoder(
             intermediate_dim=4,
             num_heads=2,
         )
@@ -126,23 +91,3 @@ class TransformerEncoderTest(TestCase):
         inputs._keras_mask = mask
         outputs = encoder(inputs)
         self.assertAllEqual(outputs._keras_mask, mask)
-
-    def test_saved_model(self):
-        model = keras.Sequential(
-            [
-                keras.Input(shape=(4, 6)),
-                transformer_encoder.TransformerEncoder(
-                    intermediate_dim=4,
-                    num_heads=2,
-                    normalize_first=True,
-                ),
-            ]
-        )
-        data = ops.random.uniform(shape=[2, 4, 6])
-        model_output = model(data)
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(path, save_format="keras_v3")
-
-        loaded_model = keras.models.load_model(path)
-        loaded_model_output = loaded_model(data)
-        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/tests/test_case.py b/keras_nlp/tests/test_case.py
index 67721ff8..a53c29a6 100644
--- a/keras_nlp/tests/test_case.py
+++ b/keras_nlp/tests/test_case.py
@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
 import tensorflow as tf
 import tree
 from absl.testing import parameterized
+from keras_core.src.backend import standardize_dtype
 
+from keras_nlp.backend import config
+from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 
 
@@ -61,3 +66,135 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         x1 = tree.map_structure(convert_to_comparible_type, x1)
         x2 = tree.map_structure(convert_to_comparible_type, x2)
         super().assertAllEqual(x1, x2, msg=msg)
+
+    def run_layer_test(
+        self,
+        layer_cls,
+        init_kwargs,
+        input_data,
+        expected_output_shape,
+        expected_output_data=None,
+        expected_num_trainable_weights=0,
+        expected_num_non_trainable_weights=0,
+        expected_num_non_trainable_variables=0,
+        run_training_check=True,
+    ):
+        # Serialization test.
+        layer = layer_cls(**init_kwargs)
+        self.run_class_serialization_test(layer)
+
+        def run_build_asserts(layer):
+            self.assertTrue(layer.built)
+            self.assertLen(
+                layer.trainable_weights,
+                expected_num_trainable_weights,
+                msg="Unexpected number of trainable_weights",
+            )
+            self.assertLen(
+                layer.non_trainable_weights,
+                expected_num_non_trainable_weights,
+                msg="Unexpected number of non_trainable_weights",
+            )
+            self.assertLen(
+                layer.non_trainable_variables,
+                expected_num_non_trainable_variables,
+                msg="Unexpected number of non_trainable_variables",
+            )
+
+        def run_output_asserts(layer, output, eager=False):
+            output_shape = tree.map_structure(
+                lambda x: None if x is None else x.shape, output
+            )
+            self.assertEqual(
+                expected_output_shape,
+                output_shape,
+                msg="Unexpected output shape",
+            )
+            output_dtype = tree.flatten(output)[0].dtype
+            self.assertEqual(
+                standardize_dtype(layer.dtype),
+                standardize_dtype(output_dtype),
+                msg="Unexpected output dtype",
+            )
+            if eager and expected_output_data is not None:
+                self.assertAllClose(expected_output_data, output)
+
+        def run_training_step(layer, input_data, output_data):
+            class TestModel(keras.Model):
+                def __init__(self, layer):
+                    super().__init__()
+                    self.layer = layer
+
+                def call(self, x):
+                    if isinstance(x, dict):
+                        return self.layer(**x)
+                    else:
+                        return self.layer(x)
+
+            model = TestModel(layer)
+            model.compile(optimizer="sgd", loss="mse", jit_compile=True)
+            model.fit(input_data, output_data, verbose=0)
+
+        if config.multi_backend():
+            # Build test.
+            layer = layer_cls(**init_kwargs)
+            if isinstance(input_data, dict):
+                shapes = {k + "_shape": v.shape for k, v in input_data.items()}
+                layer.build(**shapes)
+            else:
+                layer.build(input_data.shape)
+            run_build_asserts(layer)
+
+            # Symbolic call test.
+            keras_tensor_inputs = tree.map_structure(
+                lambda x: keras.KerasTensor(x.shape, x.dtype), input_data
+            )
+            layer = layer_cls(**init_kwargs)
+            if isinstance(keras_tensor_inputs, dict):
+                keras_tensor_outputs = layer(**keras_tensor_inputs)
+            else:
+                keras_tensor_outputs = layer(keras_tensor_inputs)
+            run_build_asserts(layer)
+            run_output_asserts(layer, keras_tensor_outputs)
+
+        # Eager call test and compiled training test.
+        layer = layer_cls(**init_kwargs)
+        if isinstance(input_data, dict):
+            output_data = layer(**input_data)
+        else:
+            output_data = layer(input_data)
+        run_output_asserts(layer, output_data, eager=True)
+
+        if run_training_check:
+            run_training_step(layer, input_data, output_data)
+
+    def run_class_serialization_test(self, instance):
+        # get_config roundtrip
+        cls = instance.__class__
+        cfg = instance.get_config()
+        cfg_json = json.dumps(cfg, sort_keys=True, indent=4)
+        ref_dir = dir(instance)[:]
+        revived_instance = cls.from_config(cfg)
+        revived_cfg = revived_instance.get_config()
+        revived_cfg_json = json.dumps(revived_cfg, sort_keys=True, indent=4)
+        self.assertEqual(cfg_json, revived_cfg_json)
+        # Dir tests only work on keras-core.
+        if config.multi_backend():
+            self.assertEqual(ref_dir, dir(revived_instance))
+
+        # serialization roundtrip
+        serialized = keras.saving.serialize_keras_object(instance)
+        serialized_json = json.dumps(serialized, sort_keys=True, indent=4)
+        revived_instance = keras.saving.deserialize_keras_object(
+            json.loads(serialized_json)
+        )
+        revived_cfg = revived_instance.get_config()
+        revived_cfg_json = json.dumps(revived_cfg, sort_keys=True, indent=4)
+        self.assertEqual(cfg_json, revived_cfg_json)
+        # Dir tests only work on keras-core.
+        if config.multi_backend():
+            new_dir = dir(revived_instance)[:]
+            for lst in [ref_dir, new_dir]:
+                if "__annotations__" in lst:
+                    lst.remove("__annotations__")
+            self.assertEqual(ref_dir, new_dir)