Repository: keras-team/keras-nlp
Pull request #1210 (2023-08-15 19:48:33)

Remove windows line endings
Randomly noticed, we have a few files that have windows line endings for some reason. Let's stick to "\n".


Patch:
diff --git a/keras_nlp/layers/modeling/position_embedding.py b/keras_nlp/layers/modeling/position_embedding.py
index 810d52d5..3b4dc638 100644
--- a/keras_nlp/layers/modeling/position_embedding.py
+++ b/keras_nlp/layers/modeling/position_embedding.py
@@ -36,6 +36,15 @@ class PositionEmbedding(keras.layers.Layer):
             to `"glorot_uniform"`.
         seq_axis: The axis of the input tensor where we add the embeddings.
 
+    Call args:
+        inputs: The tensor inputs to compute an embedding for, with shape
+            `(batch_size, sequence_length, hidden_dim)`. Only the input shape
+            will be used, as the position embedding does not depend on the
+            input sequence content.
+        start_index: An integer or integer tensor. The starting position to
+            compute the position embedding from. This is useful during cached
+            decoding, where each position is predicted separately in a loop.
+
     Examples:
 
     Called directly on input.
diff --git a/keras_nlp/layers/modeling/rotary_embedding.py b/keras_nlp/layers/modeling/rotary_embedding.py
index a4181232..e5dcff1a 100644
--- a/keras_nlp/layers/modeling/rotary_embedding.py
+++ b/keras_nlp/layers/modeling/rotary_embedding.py
@@ -38,6 +38,14 @@ class RotaryEmbedding(keras.layers.Layer):
         sequence_axis: int. Sequence axis in the input tensor.
         feature_axis: int. Feature axis in the input tensor.
 
+    Call args:
+        inputs: The tensor inputs to apply the embedding to. This can have
+            any shape, but must contain both a sequence and feature axis. The
+            rotary embedding will be applied to `inputs` and returned.
+        start_index: An integer or integer tensor. The starting position to
+            compute the rotary embedding from. This is useful during cached
+            decoding, where each position is predicted separately in a loop.
+
     Examples:
 
     ```python
@@ -74,9 +82,11 @@ class RotaryEmbedding(keras.layers.Layer):
         self.feature_axis = feature_axis
         self.scaling_factor = scaling_factor
 
-    def call(self, inputs):
+    def call(self, inputs, start_index=0):
         rotary_dim = ops.shape(inputs)[-1]
-        cos_emb, sin_emb = self._compute_cos_sin_embedding(inputs, rotary_dim)
+        cos_emb, sin_emb = self._compute_cos_sin_embedding(
+            inputs, rotary_dim, start_index
+        )
         return self._apply_rotary_pos_emb(inputs, cos_emb, sin_emb)
 
     def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
@@ -84,7 +94,7 @@ class RotaryEmbedding(keras.layers.Layer):
         half_rot_tensor = ops.concatenate((-x2, x1), axis=self.feature_axis)
         return (tensor * cos_emb) + (half_rot_tensor * sin_emb)
 
-    def _compute_cos_sin_embedding(self, x, rotary_dim):
+    def _compute_cos_sin_embedding(self, x, rotary_dim, start_index):
         freq_range = ops.arange(0, rotary_dim, 2, dtype="float32")
         freq_range = ops.cast(freq_range, self.compute_dtype)
         freq_range = freq_range / ops.cast(
@@ -95,7 +105,7 @@ class RotaryEmbedding(keras.layers.Layer):
             ** (freq_range / ops.cast(rotary_dim, self.compute_dtype))
         )
         seq_len = ops.shape(x)[self.sequence_axis]
-        tensor = ops.arange(seq_len, dtype="float32")
+        tensor = ops.arange(seq_len, dtype="float32") + start_index
         tensor = ops.cast(tensor, dtype=inverse_freq.dtype)
         freq = ops.einsum("i, j -> ij", tensor, inverse_freq)
         embedding = ops.concatenate((freq, freq), axis=self.feature_axis)
diff --git a/keras_nlp/layers/modeling/sine_position_encoding.py b/keras_nlp/layers/modeling/sine_position_encoding.py
index 70c9aee7..8d39e245 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding.py
@@ -37,6 +37,13 @@ class SinePositionEncoding(keras.layers.Layer):
             curves, as described in Attention is All You Need. Defaults to
             `10000`.
 
+    Call args:
+        inputs: The tensor inputs to compute an embedding for, with shape
+            `(batch_size, sequence_length, hidden_dim)`.
+        start_index: An integer or integer tensor. The starting position to
+            compute the encoding from. This is useful during cached decoding,
+            where each position is predicted separately in a loop.
+
     Examples:
     ```python
     # create a simple embedding layer with sinusoidal positional encoding
@@ -63,18 +70,19 @@ class SinePositionEncoding(keras.layers.Layer):
         super().__init__(**kwargs)
         self.max_wavelength = max_wavelength
 
-    def call(self, inputs):
+    def call(self, inputs, start_index=0):
         shape = ops.shape(inputs)
         seq_length = shape[-2]
         hidden_size = shape[-1]
-        position = ops.cast(ops.arange(seq_length), self.compute_dtype)
+        positions = ops.arange(seq_length)
+        positions = ops.cast(positions + start_index, self.compute_dtype)
         min_freq = ops.cast(1 / self.max_wavelength, dtype=self.compute_dtype)
         timescales = ops.power(
             min_freq,
             ops.cast(2 * (ops.arange(hidden_size) // 2), self.compute_dtype)
             / ops.cast(hidden_size, self.compute_dtype),
         )
-        angles = ops.expand_dims(position, 1) * ops.expand_dims(timescales, 0)
+        angles = ops.expand_dims(positions, 1) * ops.expand_dims(timescales, 0)
         # even indices are sine, odd are cosine
         cos_mask = ops.cast(ops.arange(hidden_size) % 2, self.compute_dtype)
         sin_mask = 1 - cos_mask
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index 3bb81209..1a84a83f 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -1,168 +1,168 @@
-# Copyright 2022 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""FNet classification model."""
-
-import copy
-
-from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
-from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
-from keras_nlp.models.f_net.f_net_backbone import f_net_kernel_initializer
-from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
-from keras_nlp.models.f_net.f_net_presets import backbone_presets
-from keras_nlp.models.task import Task
-from keras_nlp.utils.python_utils import classproperty
-
-
-@keras_nlp_export("keras_nlp.models.FNetClassifier")
-class FNetClassifier(Task):
-    """An end-to-end f_net model for classification tasks.
-
-    This model attaches a classification head to a
-    `keras_nlp.model.FNetBackbone` instance, mapping from the backbone outputs
-    to logits suitable for a classification task. For usage of this model with
-    pre-trained weights, use the `from_preset()` constructor.
-
-    This model can optionally be configured with a `preprocessor` layer, in
-    which case it will automatically apply preprocessing to raw inputs during
-    `fit()`, `predict()`, and `evaluate()`. This is done by default when
-    creating the model with `from_preset()`.
-
-    Disclaimer: Pre-trained models are provided on an "as is" basis, without
-    warranties or conditions of any kind.
-
-    Args:
-        backbone: A `keras_nlp.models.FNetBackbone` instance.
-        num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_nlp.models.FNetPreprocessor` or `None`. If
-            `None`, this model will not apply preprocessing, and inputs should
-            be preprocessed before calling the model.
-        activation: Optional `str` or callable. The
-            activation function to use on the model outputs. Set
-            `activation="softmax"` to return output probabilities.
-            Defaults to `None`.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. The dropout probability value, applied after the dense
-            layer.
-
-    Examples:
-
-    Raw string data.
-    ```python
-    features = ["The quick brown fox jumped.", "I forgot my homework."]
-    labels = [0, 3]
-
-    # Pretrained classifier.
-    classifier = keras_nlp.models.FNetClassifier.from_preset(
-        "f_net_base_en",
-        num_classes=4,
-    )
-    classifier.fit(x=features, y=labels, batch_size=2)
-    classifier.predict(x=features, batch_size=2)
-
-    # Re-compile (e.g., with a new learning rate).
-    classifier.compile(
-        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-        optimizer=keras.optimizers.Adam(5e-5),
-        jit_compile=True,
-    )
-    # Access backbone programatically (e.g., to change `trainable`).
-    classifier.backbone.trainable = False
-    # Fit again.
-    classifier.fit(x=features, y=labels, batch_size=2)
-    ```
-
-    Preprocessed integer data.
-    ```python
-    features = {
-        "token_ids": np.ones(shape=(2, 12), dtype="int32"),
-        "segment_ids": np.array([[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2),
-    }
-    labels = [0, 3]
-
-    # Pretrained classifier without preprocessing.
-    classifier = keras_nlp.models.FNetClassifier.from_preset(
-        "f_net_base_en",
-        num_classes=4,
-        preprocessor=None,
-    )
-    classifier.fit(x=features, y=labels, batch_size=2)
-    ```
-    """
-
-    def __init__(
-        self,
-        backbone,
-        num_classes,
-        preprocessor=None,
-        activation=None,
-        dropout=0.1,
-        **kwargs,
-    ):
-        inputs = backbone.input
-        pooled = backbone(inputs)["pooled_output"]
-        pooled = keras.layers.Dropout(dropout)(pooled)
-        outputs = keras.layers.Dense(
-            num_classes,
-            kernel_initializer=f_net_kernel_initializer(),
-            activation=activation,
-            name="logits",
-        )(pooled)
-        # Instantiate using Functional API Model constructor
-        super().__init__(
-            inputs=inputs,
-            outputs=outputs,
-            include_preprocessing=preprocessor is not None,
-            **kwargs,
-        )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-        self.num_classes = num_classes
-        self.activation = keras.activations.get(activation)
-        self.dropout = dropout
-
-        logit_output = self.activation == keras.activations.linear
-        self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(
-                from_logits=logit_output
-            ),
-            optimizer=keras.optimizers.Adam(5e-5),
-            metrics=[keras.metrics.SparseCategoricalAccuracy()],
-            jit_compile=True,
-        )
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "num_classes": self.num_classes,
-                "dropout": self.dropout,
-                "activation": keras.activations.serialize(self.activation),
-            }
-        )
-        return config
-
-    @classproperty
-    def backbone_cls(cls):
-        return FNetBackbone
-
-    @classproperty
-    def preprocessor_cls(cls):
-        return FNetPreprocessor
-
-    @classproperty
-    def presets(cls):
-        return copy.deepcopy(backbone_presets)
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FNet classification model."""
+
+import copy
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
+from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
+from keras_nlp.models.f_net.f_net_backbone import f_net_kernel_initializer
+from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_nlp.models.f_net.f_net_presets import backbone_presets
+from keras_nlp.models.task import Task
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras_nlp_export("keras_nlp.models.FNetClassifier")
+class FNetClassifier(Task):
+    """An end-to-end f_net model for classification tasks.
+
+    This model attaches a classification head to a
+    `keras_nlp.model.FNetBackbone` instance, mapping from the backbone outputs
+    to logits suitable for a classification task. For usage of this model with
+    pre-trained weights, use the `from_preset()` constructor.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:
+        backbone: A `keras_nlp.models.FNetBackbone` instance.
+        num_classes: int. Number of classes to predict.
+        preprocessor: A `keras_nlp.models.FNetPreprocessor` or `None`. If
+            `None`, this model will not apply preprocessing, and inputs should
+            be preprocessed before calling the model.
+        activation: Optional `str` or callable. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+            Defaults to `None`.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
+
+    Examples:
+
+    Raw string data.
+    ```python
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    labels = [0, 3]
+
+    # Pretrained classifier.
+    classifier = keras_nlp.models.FNetClassifier.from_preset(
+        "f_net_base_en",
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    classifier.predict(x=features, batch_size=2)
+
+    # Re-compile (e.g., with a new learning rate).
+    classifier.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+        jit_compile=True,
+    )
+    # Access backbone programatically (e.g., to change `trainable`).
+    classifier.backbone.trainable = False
+    # Fit again.
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Preprocessed integer data.
+    ```python
+    features = {
+        "token_ids": np.ones(shape=(2, 12), dtype="int32"),
+        "segment_ids": np.array([[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2),
+    }
+    labels = [0, 3]
+
+    # Pretrained classifier without preprocessing.
+    classifier = keras_nlp.models.FNetClassifier.from_preset(
+        "f_net_base_en",
+        num_classes=4,
+        preprocessor=None,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+    """
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        activation=None,
+        dropout=0.1,
+        **kwargs,
+    ):
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        pooled = keras.layers.Dropout(dropout)(pooled)
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=f_net_kernel_initializer(),
+            activation=activation,
+            name="logits",
+        )(pooled)
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            include_preprocessing=preprocessor is not None,
+            **kwargs,
+        )
+        # All references to `self` below this line
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.num_classes = num_classes
+        self.activation = keras.activations.get(activation)
+        self.dropout = dropout
+
+        logit_output = self.activation == keras.activations.linear
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=logit_output
+            ),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=[keras.metrics.SparseCategoricalAccuracy()],
+            jit_compile=True,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "dropout": self.dropout,
+                "activation": keras.activations.serialize(self.activation),
+            }
+        )
+        return config
+
+    @classproperty
+    def backbone_cls(cls):
+        return FNetBackbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return FNetPreprocessor
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)


Test patch:
diff --git a/keras_nlp/layers/modeling/position_embedding_test.py b/keras_nlp/layers/modeling/position_embedding_test.py
index 089aea9b..a953d814 100644
--- a/keras_nlp/layers/modeling/position_embedding_test.py
+++ b/keras_nlp/layers/modeling/position_embedding_test.py
@@ -19,7 +19,7 @@ import numpy as np
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import position_embedding
+from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
 from keras_nlp.tests.test_case import TestCase
 
 
@@ -35,9 +35,7 @@ class PositionEmbeddingTest(TestCase):
         # Create a 3-dimensional input (the first dimension is implicit).
         sequence_length = 21
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
-            sequence_length=sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=sequence_length)
         input_tensor = keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
 
@@ -52,9 +50,7 @@ class PositionEmbeddingTest(TestCase):
         # Create a 4-dimensional input (the first dimension is implicit).
         sequence_length = 21
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
-            sequence_length=sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=sequence_length)
         input_tensor = keras.Input(
             shape=(feature_size, sequence_length, feature_size)
         )
@@ -76,7 +72,7 @@ class PositionEmbeddingTest(TestCase):
         # Create a 3-dimensional input (the first dimension is implicit).
         sequence_length = 21
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
+        test_layer = PositionEmbedding(
             sequence_length=sequence_length, dtype="float16"
         )
         input_tensor = keras.Input(shape=(sequence_length, feature_size))
@@ -92,9 +88,7 @@ class PositionEmbeddingTest(TestCase):
     def test_dynamic_layer_output_shape(self):
         max_sequence_length = 40
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = keras.Input(shape=(None, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -108,9 +102,7 @@ class PositionEmbeddingTest(TestCase):
     def test_more_than_3_dimensions_dynamic(self):
         max_sequence_length = 60
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
         # Create a 4-dimensional input (the first dimension is implicit).
         input_tensor = keras.Input(shape=(None, None, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -124,9 +116,7 @@ class PositionEmbeddingTest(TestCase):
     def test_dynamic_layer_slicing(self):
         max_sequence_length = 40
         feature_size = 30
-        test_layer = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = keras.Input(shape=(None, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -149,7 +139,7 @@ class PositionEmbeddingTest(TestCase):
     def test_callable_initializer(self):
         max_sequence_length = 4
         feature_size = 3
-        test_layer = position_embedding.PositionEmbedding(
+        test_layer = PositionEmbedding(
             sequence_length=max_sequence_length,
             initializer=custom_init,
         )
@@ -170,13 +160,24 @@ class PositionEmbeddingTest(TestCase):
         )
         self.assertAllClose(model_output, expected_output)
 
+    def test_start_index(self):
+        batch_size, seq_length, feature_size = 2, 3, 4
+        layer = PositionEmbedding(seq_length)
+        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        full_output = layer(data)
+        sequential_output = ops.zeros((batch_size, seq_length, feature_size))
+        for i in range(seq_length):
+            parial_output = layer(data[:, i : i + 1, :], start_index=i)
+            sequential_output = ops.slice_update(
+                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 = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
-        )
+        test_layer = PositionEmbedding(sequence_length=max_sequence_length)
         outputs = test_layer(inputs)
         model = keras.Model(inputs=inputs, outputs=outputs)
 
@@ -196,20 +197,18 @@ class PositionEmbeddingTest(TestCase):
 
     def test_get_config_and_from_config(self):
         max_sequence_length = 40
-        test_layer = position_embedding.PositionEmbedding(
+        test_layer = PositionEmbedding(
             sequence_length=max_sequence_length,
             initializer="zeros",
         )
         config = test_layer.get_config()
-        restored = position_embedding.PositionEmbedding.from_config(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 = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
-        )
+        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)
diff --git a/keras_nlp/layers/modeling/rotary_embedding_test.py b/keras_nlp/layers/modeling/rotary_embedding_test.py
index e4b0b4cd..577d4a24 100644
--- a/keras_nlp/layers/modeling/rotary_embedding_test.py
+++ b/keras_nlp/layers/modeling/rotary_embedding_test.py
@@ -82,6 +82,19 @@ class RotaryEmbeddingTest(TestCase):
         self.assertAllClose(output[0, 0, :], expected_0, atol=0.01, rtol=0.01)
         self.assertAllClose(output[0, 3, :], expected_3, atol=0.01, rtol=0.01)
 
+    def test_start_index(self):
+        batch_size, seq_length, feature_size = 2, 3, 4
+        layer = RotaryEmbedding(seq_length)
+        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        full_output = layer(data)
+        sequential_output = ops.zeros((batch_size, seq_length, feature_size))
+        for i in range(seq_length):
+            parial_output = layer(data[:, i : i + 1, :], start_index=i)
+            sequential_output = ops.slice_update(
+                sequential_output, (0, i, 0), parial_output
+            )
+        self.assertAllClose(full_output, sequential_output)
+
     def test_get_config_and_from_config(self):
         embedding_layer = RotaryEmbedding(
             max_wavelength=1000,
diff --git a/keras_nlp/layers/modeling/sine_position_encoding_test.py b/keras_nlp/layers/modeling/sine_position_encoding_test.py
index 74119da0..7bcc6653 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding_test.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding_test.py
@@ -15,13 +15,15 @@
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import sine_position_encoding
+from keras_nlp.layers.modeling.sine_position_encoding import (
+    SinePositionEncoding,
+)
 from keras_nlp.tests.test_case import TestCase
 
 
 class SinePositionEncodingTest(TestCase):
     def test_valid_call(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        pos_encoding = SinePositionEncoding()
         model = keras.Sequential(
             [
                 keras.Input(shape=(4, 6)),
@@ -32,7 +34,7 @@ class SinePositionEncodingTest(TestCase):
         model(input)
 
     def test_static_layer_output_shape(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        pos_encoding = SinePositionEncoding()
         seq_length = 100
         hidden_size = 32
         inputs = keras.Input(shape=(seq_length, hidden_size))
@@ -44,7 +46,7 @@ class SinePositionEncodingTest(TestCase):
         self.assertEqual(expected_output_shape, outputs.shape)
 
     def test_dynamic_layer_output_shape(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        pos_encoding = SinePositionEncoding()
         hidden_size = 32
         inputs = keras.Input(shape=(None, hidden_size))
         outputs = pos_encoding(inputs)
@@ -56,7 +58,7 @@ class SinePositionEncodingTest(TestCase):
 
     # do multi dimension before sequence length
     def test_multi_dimension_layer_output_shape(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        pos_encoding = SinePositionEncoding()
         seq_length = 100
         hidden_size = 32
         inputs = keras.Input(shape=(None, seq_length, hidden_size))
@@ -68,7 +70,7 @@ class SinePositionEncodingTest(TestCase):
         self.assertEqual(expected_output_shape, outputs.shape)
 
     def test_output_correct_values(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding()
+        pos_encoding = SinePositionEncoding()
         model = keras.Sequential(
             [
                 keras.Input(shape=(4, 6)),
@@ -84,8 +86,21 @@ class SinePositionEncodingTest(TestCase):
         self.assertAllClose(output[0, 0, :], expected_0, atol=0.01, rtol=0.01)
         self.assertAllClose(output[0, 3, :], expected_3, atol=0.01, rtol=0.01)
 
+    def test_start_index(self):
+        batch_size, seq_length, feature_size = 2, 3, 4
+        layer = SinePositionEncoding()
+        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        full_output = layer(data)
+        sequential_output = ops.zeros((batch_size, seq_length, feature_size))
+        for i in range(seq_length):
+            parial_output = layer(data[:, i : i + 1, :], start_index=i)
+            sequential_output = ops.slice_update(
+                sequential_output, (0, i, 0), parial_output
+            )
+        self.assertAllClose(full_output, sequential_output)
+
     def test_get_config_and_from_config(self):
-        pos_encoding = sine_position_encoding.SinePositionEncoding(
+        pos_encoding = SinePositionEncoding(
             max_wavelength=1000,
         )
         config = pos_encoding.get_config()
@@ -93,18 +108,14 @@ class SinePositionEncodingTest(TestCase):
             "max_wavelength": 1000,
         }
         self.assertEqual(config, {**config, **expected_config_subset})
-        restored_pos_encoding = (
-            sine_position_encoding.SinePositionEncoding.from_config(config)
-        )
+        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 = sine_position_encoding.SinePositionEncoding(
-            dtype="float16"
-        )
+        pos_encoding = SinePositionEncoding(dtype="float16")
         seq_length = 100
         hidden_size = 32
         inputs = keras.Input(shape=(seq_length, hidden_size))
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
index 50392f21..614ee153 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
@@ -1,111 +1,111 @@
-# Copyright 2022 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for DistilBERT masked language model preprocessor layer."""
-
-import tensorflow as tf
-
-from keras_nlp.backend import keras
-from keras_nlp.models.distil_bert.distil_bert_masked_lm_preprocessor import (
-    DistilBertMaskedLMPreprocessor,
-)
-from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
-    DistilBertTokenizer,
-)
-from keras_nlp.tests.test_case import TestCase
-
-
-class DistilBertMaskedLMPreprocessorTest(TestCase):
-    def setUp(self):
-        self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
-        self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
-        self.vocab += ["the", "quick", "brown", "fox"]
-
-        self.preprocessor = DistilBertMaskedLMPreprocessor(
-            tokenizer=DistilBertTokenizer(
-                vocabulary=self.vocab,
-            ),
-            # Simplify our testing by masking every available token.
-            mask_selection_rate=1.0,
-            mask_token_rate=1.0,
-            random_token_rate=0.0,
-            mask_selection_length=5,
-            sequence_length=8,
-        )
-
-    def test_preprocess_strings(self):
-        input_data = " THE QUICK BROWN FOX."
-
-        x, y, sw = self.preprocessor(input_data)
-        self.assertAllEqual(x["token_ids"], [2, 4, 4, 4, 4, 4, 3, 0])
-        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
-        self.assertAllEqual(x["mask_positions"], [1, 2, 3, 4, 5])
-        self.assertAllEqual(y, [5, 6, 7, 8, 1])
-        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0, 1.0])
-
-    def test_preprocess_list_of_strings(self):
-        input_data = [" THE QUICK BROWN FOX."] * 4
-
-        x, y, sw = self.preprocessor(input_data)
-        self.assertAllEqual(x["token_ids"], [[2, 4, 4, 4, 4, 4, 3, 0]] * 4)
-        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4)
-        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4, 5]] * 4)
-        self.assertAllEqual(y, [[5, 6, 7, 8, 1]] * 4)
-        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0, 1.0]] * 4)
-
-    def test_preprocess_dataset(self):
-        sentences = tf.constant([" THE QUICK BROWN FOX."] * 4)
-        ds = tf.data.Dataset.from_tensor_slices(sentences)
-        ds = ds.map(self.preprocessor)
-        x, y, sw = ds.batch(4).take(1).get_single_element()
-        self.assertAllEqual(x["token_ids"], [[2, 4, 4, 4, 4, 4, 3, 0]] * 4)
-        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4)
-        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4, 5]] * 4)
-        self.assertAllEqual(y, [[5, 6, 7, 8, 1]] * 4)
-        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0, 1.0]] * 4)
-
-    def test_mask_multiple_sentences(self):
-        sentence_one = tf.constant(" THE QUICK")
-        sentence_two = tf.constant(" BROWN FOX.")
-
-        x, y, sw = self.preprocessor((sentence_one, sentence_two))
-        self.assertAllEqual(x["token_ids"], [2, 4, 4, 3, 4, 4, 4, 3])
-        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1])
-        self.assertAllEqual(x["mask_positions"], [1, 2, 4, 5, 6])
-        self.assertAllEqual(y, [5, 6, 7, 8, 1])
-        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0, 1.0])
-
-    def test_no_masking_zero_rate(self):
-        no_mask_preprocessor = DistilBertMaskedLMPreprocessor(
-            self.preprocessor.tokenizer,
-            mask_selection_rate=0.0,
-            mask_selection_length=5,
-            sequence_length=8,
-        )
-        input_data = " THE QUICK BROWN FOX."
-
-        x, y, sw = no_mask_preprocessor(input_data)
-        self.assertAllEqual(x["token_ids"], [2, 5, 6, 7, 8, 1, 3, 0])
-        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
-        self.assertAllEqual(x["mask_positions"], [0, 0, 0, 0, 0])
-        self.assertAllEqual(y, [0, 0, 0, 0, 0])
-        self.assertAllEqual(sw, [0.0, 0.0, 0.0, 0.0, 0.0])
-
-    def test_serialization(self):
-        config = keras.saving.serialize_keras_object(self.preprocessor)
-        new_preprocessor = keras.saving.deserialize_keras_object(config)
-        self.assertEqual(
-            new_preprocessor.get_config(),
-            self.preprocessor.get_config(),
-        )
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for DistilBERT masked language model preprocessor layer."""
+
+import tensorflow as tf
+
+from keras_nlp.backend import keras
+from keras_nlp.models.distil_bert.distil_bert_masked_lm_preprocessor import (
+    DistilBertMaskedLMPreprocessor,
+)
+from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
+    DistilBertTokenizer,
+)
+from keras_nlp.tests.test_case import TestCase
+
+
+class DistilBertMaskedLMPreprocessorTest(TestCase):
+    def setUp(self):
+        self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
+        self.vocab += ["the", "quick", "brown", "fox"]
+
+        self.preprocessor = DistilBertMaskedLMPreprocessor(
+            tokenizer=DistilBertTokenizer(
+                vocabulary=self.vocab,
+            ),
+            # Simplify our testing by masking every available token.
+            mask_selection_rate=1.0,
+            mask_token_rate=1.0,
+            random_token_rate=0.0,
+            mask_selection_length=5,
+            sequence_length=8,
+        )
+
+    def test_preprocess_strings(self):
+        input_data = " THE QUICK BROWN FOX."
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [2, 4, 4, 4, 4, 4, 3, 0])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
+        self.assertAllEqual(x["mask_positions"], [1, 2, 3, 4, 5])
+        self.assertAllEqual(y, [5, 6, 7, 8, 1])
+        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0, 1.0])
+
+    def test_preprocess_list_of_strings(self):
+        input_data = [" THE QUICK BROWN FOX."] * 4
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [[2, 4, 4, 4, 4, 4, 3, 0]] * 4)
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4)
+        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4, 5]] * 4)
+        self.assertAllEqual(y, [[5, 6, 7, 8, 1]] * 4)
+        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0, 1.0]] * 4)
+
+    def test_preprocess_dataset(self):
+        sentences = tf.constant([" THE QUICK BROWN FOX."] * 4)
+        ds = tf.data.Dataset.from_tensor_slices(sentences)
+        ds = ds.map(self.preprocessor)
+        x, y, sw = ds.batch(4).take(1).get_single_element()
+        self.assertAllEqual(x["token_ids"], [[2, 4, 4, 4, 4, 4, 3, 0]] * 4)
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4)
+        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4, 5]] * 4)
+        self.assertAllEqual(y, [[5, 6, 7, 8, 1]] * 4)
+        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0, 1.0]] * 4)
+
+    def test_mask_multiple_sentences(self):
+        sentence_one = tf.constant(" THE QUICK")
+        sentence_two = tf.constant(" BROWN FOX.")
+
+        x, y, sw = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(x["token_ids"], [2, 4, 4, 3, 4, 4, 4, 3])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1])
+        self.assertAllEqual(x["mask_positions"], [1, 2, 4, 5, 6])
+        self.assertAllEqual(y, [5, 6, 7, 8, 1])
+        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0, 1.0])
+
+    def test_no_masking_zero_rate(self):
+        no_mask_preprocessor = DistilBertMaskedLMPreprocessor(
+            self.preprocessor.tokenizer,
+            mask_selection_rate=0.0,
+            mask_selection_length=5,
+            sequence_length=8,
+        )
+        input_data = " THE QUICK BROWN FOX."
+
+        x, y, sw = no_mask_preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [2, 5, 6, 7, 8, 1, 3, 0])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
+        self.assertAllEqual(x["mask_positions"], [0, 0, 0, 0, 0])
+        self.assertAllEqual(y, [0, 0, 0, 0, 0])
+        self.assertAllEqual(sw, [0.0, 0.0, 0.0, 0.0, 0.0])
+
+    def test_serialization(self):
+        config = keras.saving.serialize_keras_object(self.preprocessor)
+        new_preprocessor = keras.saving.deserialize_keras_object(config)
+        self.assertEqual(
+            new_preprocessor.get_config(),
+            self.preprocessor.get_config(),
+        )
diff --git a/keras_nlp/models/f_net/f_net_classifier_test.py b/keras_nlp/models/f_net/f_net_classifier_test.py
index 917886b2..f30e043d 100644
--- a/keras_nlp/models/f_net/f_net_classifier_test.py
+++ b/keras_nlp/models/f_net/f_net_classifier_test.py
@@ -1,148 +1,148 @@
-# Copyright 2023 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tests for FNet classification model."""
-
-import io
-import os
-
-import numpy as np
-import pytest
-import sentencepiece
-import tensorflow as tf
-
-from keras_nlp.backend import keras
-from keras_nlp.backend import ops
-from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
-from keras_nlp.models.f_net.f_net_classifier import FNetClassifier
-from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
-from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
-from keras_nlp.tests.test_case import TestCase
-
-
-class FNetClassifierTest(TestCase):
-    def setUp(self):
-        # Setup Model
-        bytes_io = io.BytesIO()
-        vocab_data = tf.data.Dataset.from_tensor_slices(
-            ["the quick brown fox", "the earth is round"]
-        )
-
-        sentencepiece.SentencePieceTrainer.train(
-            sentence_iterator=vocab_data.as_numpy_iterator(),
-            model_writer=bytes_io,
-            vocab_size=12,
-            model_type="WORD",
-            pad_id=3,
-            unk_id=0,
-            bos_id=4,
-            eos_id=5,
-            pad_piece="<pad>",
-            unk_piece="<unk>",
-            bos_piece="[CLS]",
-            eos_piece="[SEP]",
-            user_defined_symbols="[MASK]",
-        )
-
-        self.proto = bytes_io.getvalue()
-
-        self.preprocessor = FNetPreprocessor(
-            tokenizer=FNetTokenizer(proto=self.proto),
-            sequence_length=8,
-        )
-        self.backbone = FNetBackbone(
-            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
-            num_layers=2,
-            hidden_dim=2,
-            intermediate_dim=4,
-            max_sequence_length=self.preprocessor.packer.sequence_length,
-        )
-        self.classifier = FNetClassifier(
-            self.backbone,
-            num_classes=4,
-            preprocessor=self.preprocessor,
-            # Check we handle serialization correctly.
-            activation=keras.activations.softmax,
-        )
-
-        # Setup data.
-        self.raw_batch = [
-            "the quick brown fox.",
-            "the slow brown fox.",
-        ]
-        self.preprocessed_batch = self.preprocessor(self.raw_batch)
-        self.raw_dataset = tf.data.Dataset.from_tensor_slices(
-            (self.raw_batch, np.ones((2,)))
-        ).batch(2)
-        self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
-
-    def test_valid_call_classifier(self):
-        self.classifier(self.preprocessed_batch)
-
-    def test_classifier_predict(self):
-        preds1 = self.classifier.predict(self.raw_batch)
-        self.classifier.preprocessor = None
-        preds2 = self.classifier.predict(self.preprocessed_batch)
-        # Assert predictions match.
-        self.assertAllClose(preds1, preds2)
-        # Assert valid softmax output.
-        self.assertAllClose(ops.sum(preds2, axis=-1), [1.0, 1.0])
-
-    def test_fnet_classifier_fit(self):
-        self.classifier.fit(self.raw_dataset)
-        self.classifier.preprocessor = None
-        self.classifier.fit(self.preprocessed_dataset)
-
-    def test_classifier_fit_no_xla(self):
-        self.classifier.preprocessor = None
-        self.classifier.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
-            jit_compile=False,
-        )
-        self.classifier.fit(self.preprocessed_dataset)
-
-    def test_serialization(self):
-        # Defaults.
-        original = FNetClassifier(
-            self.backbone,
-            num_classes=2,
-        )
-        config = keras.saving.serialize_keras_object(original)
-        restored = keras.saving.deserialize_keras_object(config)
-        self.assertEqual(restored.get_config(), original.get_config())
-        # With options.
-        original = FNetClassifier(
-            self.backbone,
-            num_classes=4,
-            preprocessor=self.preprocessor,
-            activation=keras.activations.softmax,
-            name="test",
-            trainable=False,
-        )
-        config = keras.saving.serialize_keras_object(original)
-        restored = keras.saving.deserialize_keras_object(config)
-        self.assertEqual(restored.get_config(), original.get_config())
-
-    @pytest.mark.large
-    def test_saved_model(self):
-        model_output = self.classifier.predict(self.raw_batch)
-        path = os.path.join(self.get_temp_dir(), "model.keras")
-        self.classifier.save(path, save_format="keras_v3")
-        restored_model = keras.models.load_model(path)
-
-        # Check we got the real object back.
-        self.assertIsInstance(restored_model, FNetClassifier)
-
-        # Check that output matches.
-        restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for FNet classification model."""
+
+import io
+import os
+
+import numpy as np
+import pytest
+import sentencepiece
+import tensorflow as tf
+
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
+from keras_nlp.models.f_net.f_net_classifier import FNetClassifier
+from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class FNetClassifierTest(TestCase):
+    def setUp(self):
+        # Setup Model
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=12,
+            model_type="WORD",
+            pad_id=3,
+            unk_id=0,
+            bos_id=4,
+            eos_id=5,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
+        )
+
+        self.proto = bytes_io.getvalue()
+
+        self.preprocessor = FNetPreprocessor(
+            tokenizer=FNetTokenizer(proto=self.proto),
+            sequence_length=8,
+        )
+        self.backbone = FNetBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=self.preprocessor.packer.sequence_length,
+        )
+        self.classifier = FNetClassifier(
+            self.backbone,
+            num_classes=4,
+            preprocessor=self.preprocessor,
+            # Check we handle serialization correctly.
+            activation=keras.activations.softmax,
+        )
+
+        # Setup data.
+        self.raw_batch = [
+            "the quick brown fox.",
+            "the slow brown fox.",
+        ]
+        self.preprocessed_batch = self.preprocessor(self.raw_batch)
+        self.raw_dataset = tf.data.Dataset.from_tensor_slices(
+            (self.raw_batch, np.ones((2,)))
+        ).batch(2)
+        self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
+
+    def test_valid_call_classifier(self):
+        self.classifier(self.preprocessed_batch)
+
+    def test_classifier_predict(self):
+        preds1 = self.classifier.predict(self.raw_batch)
+        self.classifier.preprocessor = None
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(ops.sum(preds2, axis=-1), [1.0, 1.0])
+
+    def test_fnet_classifier_fit(self):
+        self.classifier.fit(self.raw_dataset)
+        self.classifier.preprocessor = None
+        self.classifier.fit(self.preprocessed_dataset)
+
+    def test_classifier_fit_no_xla(self):
+        self.classifier.preprocessor = None
+        self.classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
+            jit_compile=False,
+        )
+        self.classifier.fit(self.preprocessed_dataset)
+
+    def test_serialization(self):
+        # Defaults.
+        original = FNetClassifier(
+            self.backbone,
+            num_classes=2,
+        )
+        config = keras.saving.serialize_keras_object(original)
+        restored = keras.saving.deserialize_keras_object(config)
+        self.assertEqual(restored.get_config(), original.get_config())
+        # With options.
+        original = FNetClassifier(
+            self.backbone,
+            num_classes=4,
+            preprocessor=self.preprocessor,
+            activation=keras.activations.softmax,
+            name="test",
+            trainable=False,
+        )
+        config = keras.saving.serialize_keras_object(original)
+        restored = keras.saving.deserialize_keras_object(config)
+        self.assertEqual(restored.get_config(), original.get_config())
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        model_output = self.classifier.predict(self.raw_batch)
+        path = os.path.join(self.get_temp_dir(), "model.keras")
+        self.classifier.save(path, save_format="keras_v3")
+        restored_model = keras.models.load_model(path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, FNetClassifier)
+
+        # Check that output matches.
+        restored_output = restored_model.predict(self.raw_batch)
+        self.assertAllClose(model_output, restored_output)