Repository: keras-team/keras-nlp
Pull request #1238 (2023-09-08 00:38:18)

Add a test harness based on keras-core's `run_layer_test`
Eventually, I want to add some dtype changes similar to -> https://github.com/keras-team/keras-core/pull/805
But the nice for thing for that PR on keras-core was I could add dtype test to a common harness and test all layers.

So I think it's finally time to bite the bullet and add these for keras-nlp. This ports and simplifies [run_layer_test](https://github.com/keras-team/keras-core/blob/ac9be3319e8a907e4db4348421a6488e258517a2/keras_core/testing/test_case.py#L108) from keras-core and applies it to our modeling layers.

I am ditching the saved model tests for our individual layers, with the idea being that saved model tests are slow, and we get fairly robust serialization tests now through the harness. If this is good enough for keras-core layers, we can follow suit here. We still test saving end to end through our `models` tests.

mattdangerw: /gcbrun

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:


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)