Repository: keras-team/keras-nlp
Pull request #146 (2022-04-28 16:04:30)

Docstring updates for upcoming doc publish


Patch:
diff --git a/examples/bert/bert_model.py b/examples/bert/bert_model.py
index 3350fdc7..99df0440 100644
--- a/examples/bert/bert_model.py
+++ b/examples/bert/bert_model.py
@@ -488,7 +488,7 @@ class BertModel(keras.Model):
 
         self._position_embedding_layer = keras_nlp.layers.PositionEmbedding(
             initializer=initializer,
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             name="position_embedding",
         )
 
diff --git a/examples/bert/run_glue_finetuning.py b/examples/bert/run_glue_finetuning.py
index 0bc7ffdf..156bc601 100644
--- a/examples/bert/run_glue_finetuning.py
+++ b/examples/bert/run_glue_finetuning.py
@@ -16,6 +16,7 @@
 import json
 
 import datasets
+import keras_tuner
 import tensorflow as tf
 import tensorflow_text as tftext
 from absl import app
@@ -68,8 +69,6 @@ flags.DEFINE_integer("batch_size", 32, "The batch size.")
 
 flags.DEFINE_integer("epochs", 3, "The number of training epochs.")
 
-flags.DEFINE_float("learning_rate", 2e-5, "The initial learning rate for Adam.")
-
 flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
 
 
@@ -167,9 +166,32 @@ class BertClassificationFinetuner(keras.Model):
         return self._logit_layer(outputs)
 
 
+class BertHyperModel(keras_tuner.HyperModel):
+    """Creates a hypermodel to help with the search space for finetuning."""
+
+    def __init__(self, bert_config):
+        self.bert_config = bert_config
+
+    def build(self, hp):
+        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
+        bert_config = self.bert_config
+        finetuning_model = BertClassificationFinetuner(
+            bert_model=model,
+            hidden_size=bert_config["hidden_size"],
+            num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
+        )
+        finetuning_model.compile(
+            optimizer=keras.optimizers.Adam(
+                learning_rate=hp.Choice("lr", [5e-5, 4e-5, 3e-5, 2e-5])
+            ),
+            loss="sparse_categorical_crossentropy",
+            metrics=["accuracy"],
+        )
+        return finetuning_model
+
+
 def main(_):
     print(f"Reading input model from {FLAGS.saved_model_input}")
-    model = keras.models.load_model(FLAGS.saved_model_input)
 
     vocab = []
     with open(FLAGS.vocab_file, "r") as vocab_file:
@@ -200,6 +222,7 @@ def main(_):
 
     # Read and preprocess GLUE task data.
     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
+
     train_ds = train_ds.batch(FLAGS.batch_size).map(
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
@@ -210,18 +233,27 @@ def main(_):
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
 
-    finetuning_model = BertClassificationFinetuner(
-        bert_model=model,
-        hidden_size=bert_config["hidden_size"],
-        num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-    )
-    finetuning_model.compile(
-        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),
-        loss="sparse_categorical_crossentropy",
-        metrics=["accuracy"],
+    # Create a hypermodel object for a RandomSearch.
+    hypermodel = BertHyperModel(bert_config)
+
+    # Initialize the random search over the 4 learning rate parameters, for 4
+    # trials and 3 epochs for each trial.
+    tuner = keras_tuner.RandomSearch(
+        hypermodel=hypermodel,
+        objective=keras_tuner.Objective("val_loss", direction="min"),
+        max_trials=4,
+        overwrite=True,
+        project_name="hyperparameter_tuner_results",
     )
-    finetuning_model.fit(
-        train_ds, epochs=FLAGS.epochs, validation_data=validation_ds
+
+    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)
+
+    # Extract the best hyperparameters after the search.
+    best_hp = tuner.get_best_hyperparameters()[0]
+    finetuning_model = tuner.get_best_models()[0]
+
+    print(
+        f"The best hyperparameters found are:\nLearning Rate: {best_hp['lr']}"
     )
 
     if FLAGS.do_evaluation:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index ba357d34..edbe34c3 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -15,5 +15,6 @@
 from keras_nlp import layers
 from keras_nlp import metrics
 from keras_nlp import tokenizers
+from keras_nlp import utils
 
 __version__ = "0.2.0-dev.1"
diff --git a/keras_nlp/layers/fnet_encoder.py b/keras_nlp/layers/fnet_encoder.py
index 47937392..90a3fc33 100644
--- a/keras_nlp/layers/fnet_encoder.py
+++ b/keras_nlp/layers/fnet_encoder.py
@@ -21,12 +21,11 @@ from tensorflow import keras
 class FNetEncoder(keras.layers.Layer):
     """FNet encoder.
 
-    This class follows the architecture of FNet encoder layer in paper
-    "FNet: Mixing Tokens with Fourier Transforms"
-    (https://arxiv.org/abs/2105.03824). Users can instantiate multiple instances
-    of this class to stack up the encoder.
+    This class follows the architecture of FNet encoder layer in the
+    [FNet paper](https://arxiv.org/abs/2105.03824). Users can instantiate
+    multiple instances of this class to stack up the encoder.
 
-    Note on padding: In the official FNet code, padding tokens are added to the
+    Note on masking: In the official FNet code, padding tokens are added to the
     the input. However, the padding masks are deleted, i.e., mixing of
     all tokens is done. This is because certain frequencies will be zeroed
     out if we apply padding masks in every encoder layer. Hence, we don't
@@ -36,14 +35,14 @@ class FNetEncoder(keras.layers.Layer):
         intermediate_dim: int. The hidden size of feedforward network.
         dropout: float, defaults to 0. The dropout value, applied in the
             feedforward network.
-        activation: string or `tf.keras.activations`, defaults to "relu". The
+        activation: string or `keras.activations`, defaults to "relu". The
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
-        kernel_initializer: "string" or `tf.keras.initializers` initializer,
+        kernel_initializer: "string" or `keras.initializers` initializer,
             defaults to "glorot_uniform". The kernel initializer for the dense
             layers.
-        bias_initializer: "string" or `tf.keras.initializers` initializer,
+        bias_initializer: "string" or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for the dense layers.
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
@@ -56,9 +55,9 @@ class FNetEncoder(keras.layers.Layer):
         intermediate_dim=64)
 
     # Create a simple model containing the encoder.
-    input = tf.keras.Input(shape=[4, 6])
+    input = keras.Input(shape=[10, 64])
     output = encoder(input)
-    model = tf.keras.Model(inputs=input, outputs=output)
+    model = keras.Model(inputs=input, outputs=output)
 
     # Call encoder on the inputs.
     input_data = tf.random.uniform(shape=[1, 10, 64])
@@ -66,7 +65,7 @@ class FNetEncoder(keras.layers.Layer):
     ```
 
     References:
-        [Lee-Thorp et al., 2021](https://arxiv.org/abs/2105.03824)
+     - [Lee-Thorp et al., 2021](https://arxiv.org/abs/2105.03824)
     """
 
     def __init__(
diff --git a/keras_nlp/layers/mlm_head.py b/keras_nlp/layers/mlm_head.py
index 09644205..53c1a7ff 100644
--- a/keras_nlp/layers/mlm_head.py
+++ b/keras_nlp/layers/mlm_head.py
@@ -22,10 +22,11 @@ class MLMHead(keras.layers.Layer):
     """Masked Language Model (MLM) head.
 
     This layer takes two inputs:
+
      - `inputs`: which should be a tensor of encoded tokens with shape
-            `(batch_size, sequence_length, encoding_dim)`.
+       `(batch_size, sequence_length, encoding_dim)`.
      - `mask_positions`: which should be a tensor of integer positions to
-            predict with shape `(batch_size, masks_per_sequence)`.
+       predict with shape `(batch_size, masks_per_sequence)`.
 
     The token encodings should usually be the last output of an encoder model,
     and mask positions should be the interger positions you would like to
@@ -46,17 +47,18 @@ class MLMHead(keras.layers.Layer):
         embedding_weights: Optional. The weights of the word embedding used
             to transform input token ids. The transpose of this weight matrix
             will be used to project a token embedding vector to a prediction
-            over all input words, as described in [1].
+            over all input words, as described
+            [here](https://arxiv.org/abs/1608.05859).
         intermediate_activation: The activation function of inner dense layer.
         activation: The activation function for the outputs of the layer.
             Usually either `None` (return logits), or `"softmax"`
             (return probabilities).
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
-        kernel_initializer: string or tf.keras.initializers initializer,
+        kernel_initializer: string or `keras.initializers` initializer,
             defaults to "glorot_uniform". The kernel initializer for
             the dense and multiheaded attention layers.
-        bias_initializer: string or tf.keras.initializers initializer,
+        bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded attention layers.
         name: string, defaults to None. The name of the layer.
@@ -91,7 +93,7 @@ class MLMHead(keras.layers.Layer):
     ```
 
     References:
-        [1] [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+     - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
     """
 
     def __init__(
@@ -147,7 +149,7 @@ class MLMHead(keras.layers.Layer):
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
         )
-        self._layer_norm = tf.keras.layers.LayerNormalization(
+        self._layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
         if self.embedding_weights is None:
diff --git a/keras_nlp/layers/mlm_mask_generator.py b/keras_nlp/layers/mlm_mask_generator.py
index 85b207a4..6576dc50 100644
--- a/keras_nlp/layers/mlm_mask_generator.py
+++ b/keras_nlp/layers/mlm_mask_generator.py
@@ -73,15 +73,15 @@ class MLMMaskGenerator(keras.layers.Layer):
     Examples:
 
     Basic usage.
-    >>> masker = keras_nlp.layers.MLMMaskGenerator( \
-            vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \
-            mask_selection_length=5)
+    >>> masker = keras_nlp.layers.MLMMaskGenerator(
+    ...     vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0,
+    ...     mask_selection_length=5)
     >>> masker(tf.constant([1, 2, 3, 4, 5]))
 
     Ragged Input:
-    >>> masker = keras_nlp.layers.MLMMaskGenerator( \
-            vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \
-            mask_selection_length=5)
+    >>> masker = keras_nlp.layers.MLMMaskGenerator(
+    ...     vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0,
+    ...     mask_selection_length=5)
     >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))
     """
 
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
index 5d3b01f8..2012f010 100644
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -32,18 +32,18 @@ class PositionEmbedding(keras.layers.Layer):
     corresponds to the sequence, that is, the penultimate dimension.
 
     Args:
-        max_length: The maximum length of the dynamic sequence.
+        sequence_length: The maximum length of the dynamic sequence.
         initializer: The initializer to use for the embedding weights. Defaults
-            to "glorot_uniform".
+            to `"glorot_uniform"`.
         seq_axis: The axis of the input tensor where we add the embeddings.
 
-    Example:
+    Examples:
     ```python
     token_embeddings = layers.Embedding(
         input_dim=vocab_size, output_dim=embed_dim
     )
     position_embeddings = keras_nlp.layers.PositionEmbedding(
-        max_length=max_length
+        sequence_length=sequence_length
     )
 
     embedded_tokens = token_embeddings(inputs)
@@ -52,29 +52,28 @@ class PositionEmbedding(keras.layers.Layer):
     ```
 
     Reference:
-        [BERT: Pre-training of Deep Bidirectional Transformers for Language
-        Understanding](https://arxiv.org/abs/1810.04805).
+     - [Devlin et al., 2019](https://arxiv.org/abs/1810.04805)
     """
 
     def __init__(
         self,
-        max_length,
+        sequence_length,
         initializer="glorot_uniform",
         **kwargs,
     ):
         super().__init__(**kwargs)
-        if max_length is None:
+        if sequence_length is None:
             raise ValueError(
-                "`max_length` must be an Integer, received `None`."
+                "`sequence_length` must be an Integer, received `None`."
             )
-        self.max_length = int(max_length)
+        self.sequence_length = int(sequence_length)
         self.initializer = keras.initializers.get(initializer)
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
-                "max_length": self.max_length,
+                "sequence_length": self.sequence_length,
                 "initializer": keras.initializers.serialize(self.initializer),
             }
         )
@@ -84,7 +83,7 @@ class PositionEmbedding(keras.layers.Layer):
         feature_size = input_shape[-1]
         self.position_embeddings = self.add_weight(
             "embeddings",
-            shape=[self.max_length, feature_size],
+            shape=[self.sequence_length, feature_size],
             initializer=self.initializer,
             trainable=True,
         )
diff --git a/keras_nlp/layers/sine_position_encoding.py b/keras_nlp/layers/sine_position_encoding.py
index 4112ca8a..0ce8ec3c 100644
--- a/keras_nlp/layers/sine_position_encoding.py
+++ b/keras_nlp/layers/sine_position_encoding.py
@@ -35,7 +35,7 @@ class SinePositionEncoding(keras.layers.Layer):
             curves, as described in Attention is All You Need. Defaults to
             10000.
 
-    Example:
+    Examples:
     ```python
     # create a simple embedding layer with sinusoidal positional encoding
     seq_len = 100
@@ -50,7 +50,7 @@ class SinePositionEncoding(keras.layers.Layer):
     ```
 
     References:
-      [Attention is All You Need](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
     """
 
     def __init__(
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index 035ddd02..c6b9fb2d 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -28,7 +28,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     Args:
         vocabulary_size: The size of the vocabulary (should be no larger
             than 999)
-        max_length: The maximum length of input sequence
+        sequence_length: The maximum length of input sequence
         embedding_dim: The output dimension of the embedding layer
         embeddings_initializer: The initializer to use for the Embedding
             Layers
@@ -41,7 +41,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             used in the vocabulary
             (input_dim should equal size of vocabulary + 1).
 
-    Example:
+    Examples:
     ```python
     seq_length = 50
     vocab_size = 5000
@@ -49,7 +49,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     inputs = tf.keras.Input(shape=(seq_length,))
     embedding_layer = keras_nlp.layers.TokenAndPositionEmbedding(
         vocabulary_size=vocab_size,
-        max_length=max_length,
+        sequence_length=seq_length,
         embedding_dim=embed_dim,
     )
     outputs = embedding_layer(inputs)
@@ -59,7 +59,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     def __init__(
         self,
         vocabulary_size,
-        max_length,
+        sequence_length,
         embedding_dim,
         embeddings_initializer="glorot_uniform",
         mask_zero=False,
@@ -70,16 +70,16 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             raise ValueError(
                 "`vocabulary_size` must be an Integer, received `None`."
             )
-        if max_length is None:
+        if sequence_length is None:
             raise ValueError(
-                "`max_length` must be an Integer, received `None`."
+                "`sequence_length` must be an Integer, received `None`."
             )
         if embedding_dim is None:
             raise ValueError(
                 "`embedding_dim` must be an Integer, received `None`."
             )
         self.vocabulary_size = int(vocabulary_size)
-        self.max_length = int(max_length)
+        self.sequence_length = int(sequence_length)
         self.embedding_dim = int(embedding_dim)
         self.token_embedding = keras.layers.Embedding(
             vocabulary_size,
@@ -88,7 +88,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             mask_zero=mask_zero,
         )
         self.position_embedding = keras_nlp.layers.PositionEmbedding(
-            max_length=max_length,
+            sequence_length=sequence_length,
             initializer=embeddings_initializer,
         )
         self.supports_masking = self.token_embedding.supports_masking
@@ -98,7 +98,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         config.update(
             {
                 "vocabulary_size": self.vocabulary_size,
-                "max_length": self.max_length,
+                "sequence_length": self.sequence_length,
                 "embedding_dim": self.embedding_dim,
                 "embeddings_initializer": keras.initializers.serialize(
                     self.token_embedding.embeddings_initializer
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index 92d59c2e..f8cdd105 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -26,23 +26,29 @@ from keras_nlp.layers.transformer_layer_utils import (  # isort:skip
 class TransformerDecoder(keras.layers.Layer):
     """Transformer decoder.
 
-    This class follows the architecture of transformer decoder layer in paper
-    "Attention is All You Need" (https://arxiv.org/abs/1706.03762). Users can
-    instantiate multiple instances of this class to stack up the decoder.
+    This class follows the architecture of the transformer decoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up a decoder.
+
+    This layer will correctly compute an attention mask from an implicit
+    Keras padding mask (for example, by passing `mask_zero=True` to a
+    `keras.layers.Embedding` layer). See the Masking and Padding
+    [guide](https://keras.io/guides/understanding_masking_and_padding/)
+    for more details.
 
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
         num_heads: int, the number of heads in MultiHeadAttention.
         dropout: float, defaults to 0. the dropout value, shared by
             MultiHeadAttention and feedforward network.
-        activation: string or tf.keras.activations, defaults to "relu". the
+        activation: string or `keras.activations`, defaults to "relu". the
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The eps value in layer
             normalization components.
-        kernel_initializer: string or tf.keras.initializers initializer,
+        kernel_initializer: string or `keras.initializers` initializer,
             defaults to "glorot_uniform". The kernel initializer for
             the dense and multiheaded attention layers.
-        bias_initializer: string or tf.keras.initializers initializer,
+        bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded attention layers.
         name: string, defaults to None. The name of the layer.
@@ -55,21 +61,21 @@ class TransformerDecoder(keras.layers.Layer):
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the decoder.
-    decoder_input = tf.keras.Input(shape=[4, 6])
-    encoder_input = tf.keras.Input(shape=[4, 6])
+    decoder_input = keras.Input(shape=[10, 64])
+    encoder_input = keras.Input(shape=[10, 64])
     output = decoder(decoder_input, encoder_input)
-    model = tf.keras.Model(inputs=[decoder_input, encoder_input],
+    model = keras.Model(inputs=[decoder_input, encoder_input],
         outputs=output)
 
     # Call decoder on the inputs.
-    decoder_input_data = tf.random.uniform(shape=[1, 10, 64])
-    encoder_input_data = tf.random.uniform(shape=[1, 10, 64])
+    decoder_input_data = tf.random.uniform(shape=[2, 10, 64])
+    encoder_input_data = tf.random.uniform(shape=[2, 10, 64])
     decoder_output = model([decoder_input_data, encoder_input_data])
 
     ```
 
     References:
-        [Vaswani et al., 20XX](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
 
     """
 
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
index 870228be..c67409d9 100644
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -24,23 +24,30 @@ from keras_nlp.layers.transformer_layer_utils import (  # isort:skip
 class TransformerEncoder(keras.layers.Layer):
     """Transformer encoder.
 
-    This class follows the architecture of transformer encoder layer in paper
-    "Attention is All You Need" (https://arxiv.org/abs/1706.03762). Users can
-    instantiate multiple instances of this class to stack up the encoder.
+    This class follows the architecture of the transformer encoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up an encoder.
+
+    This layer will correctly compute an attention mask from an implicit
+    Keras padding mask (for example, by passing `mask_zero=True` to a
+    `keras.layers.Embedding` layer). See the Masking and Padding
+    [guide](https://keras.io/guides/understanding_masking_and_padding/)
+    for more details.
 
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
-        num_heads: int, the number of heads in MultiHeadAttention.
+        num_heads: int, the number of heads in the
+            `keras.layers.MultiHeadAttention` layer.
         dropout: float, defaults to 0. the dropout value, shared by
-            MultiHeadAttention and feedforward network.
-        activation: string or `tf.keras.activations`, defaults to "relu". the
+            `keras.layers.MultiHeadAttention` and feedforward network.
+        activation: string or `keras.activations`, defaults to "relu". the
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
-        kernel_initializer: string or tf.keras.initializers initializer,
+        kernel_initializer: string or `keras.initializers` initializer,
             defaults to "glorot_uniform". The kernel initializer for
             the dense and multiheaded attention layers.
-        bias_initializer: string or tf.keras.initializers initializer,
+        bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded attention layers.
         name: string, defaults to None. The name of the layer.
@@ -54,18 +61,18 @@ class TransformerEncoder(keras.layers.Layer):
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the encoder.
-    input = tf.keras.Input(shape=[4, 6])
+    input = keras.Input(shape=[10, 64])
     output = encoder(input)
-    model = tf.keras.Model(inputs=input, outputs=output)
+    model = keras.Model(inputs=input, outputs=output)
 
     # Call encoder on the inputs.
-    input_data = tf.random.uniform(shape=[1, 10, 64])
+    input_data = tf.random.uniform(shape=[2, 10, 64])
     output = model(input_data)
 
     ```
 
     References:
-        [Vaswani et al., 20XX](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
     """
 
     def __init__(
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
index 1a106cd9..fcbad637 100644
--- a/keras_nlp/tokenizers/byte_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -30,19 +30,25 @@ class ByteTokenizer(tokenizer.Tokenizer):
     This tokenizer is a vocabulary-free tokenizer which will tokenize text as
     as raw bytes from [0, 256).
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
     If input is a batch of strings:
     By default, the layer will output a `tf.RaggedTensor` where the last
     dimension of the output is ragged. If `sequence_length` is set, the layer
     will output a dense `tf.Tensor` where all inputs have been padded or
-    truncated to `sequence_length`. The output dtype can be controlled via the
-    `dtype` argument, which should be an integer type
-    (tf.int16, tf.int32, etc.).
+    truncated to `sequence_length`.
 
     If input is a scalar string:
     There are two cases here. If `sequence_length` is set, the output will be
     a dense `tf.Tensor` of shape `[sequence_length]`. Otherwise, the output will
     be a dense `tf.Tensor` of shape `[None]`.
 
+    The output dtype can be controlled via the
+    `dtype` argument, which should be an integer type
+    (tf.int16, tf.int32, etc.).
+
     Args:
         lowercase: boolean. If True, the input text will be converted to
             lowercase before tokenization.
@@ -89,7 +95,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     <tf.Tensor: shape=(1, 8), dtype=int32, numpy=
     array([[104, 101, 108, 108, 111,   0,   0,   0]], dtype=int32)>
 
-    Tokenize first, then batch the dataset up.
+    Tokenize, then batch for ragged outputs.
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.map(tokenizer)
@@ -97,14 +103,14 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> ds.take(1).get_single_element()
     <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
 
-    Batch the inputs and then tokenize.
+    Batch, then tokenize for ragged outputs.
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.batch(2).map(tokenizer)
     >>> ds.take(1).get_single_element()
     <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
 
-    Tokenize first, then batch the dataset up (`sequence_length` provided).
+    Tokenize, then batch for dense outputs (`sequence_length` provided).
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.map(tokenizer)
@@ -114,7 +120,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     array([[104, 101, 108, 108, 111],
            [102, 117, 110,   0,   0]], dtype=int32)>
 
-    Batch the inputs and then tokenize (`sequence_length` provided).
+    Batch, then tokenize for dense outputs. (`sequence_length` provided).
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.batch(2).map(tokenizer)
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
index baccd457..03602c82 100644
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -20,32 +20,42 @@ from tensorflow import keras
 class Tokenizer(keras.layers.Layer):
     """A base class for tokenizer layers.
 
-    The class is intended as a base class when implementing a tokenizer as a
-    `keras.layers.Layer`. It contains two new methods `tokenize()` and
-    `detokenize()`.
+    Tokenizers in the KerasNLP library should all subclass this layer.
+    The class provides two core methods `tokenize()` and `detokenize()` for
+    going from plain text to sequences and back. A tokenizer is a subclass of
+    `keras.layers.Layer` and can be combined into a `keras.Model`.
 
-    Subclassers should always implement the `tokenize` method, which will also
-    be the default when invoking the layer directly on inputs.
+    Subclassers should always implement the `tokenize()` method, which will also
+    be the default when calling the layer directly on inputs.
 
-    Subclassers should implement the `detokenize()` method if the tokenization
-    is reversible. Otherwise, this can be skipped.
+    Subclassers can optionally implement the `detokenize()` method if the
+    tokenization is reversible. Otherwise, this can be skipped.
 
     Subclassers should implement `get_vocabulary()`, `vocabulary_size()`,
     `token_to_id()` and `id_to_token()` if applicable. For some simple
-    "pre-tokenizers," such whitespace splitting, these methods may not apply,
-    and can be skipped.
+    "vocab free" tokenizers, such as a whitespace splitter show below, these
+    methods do not apply and can be skipped.
 
     Examples:
 
     ```python
     class WhitespaceSplitterTokenizer(keras_nlp.tokenizers.Tokenizer):
         def tokenize(self, inputs):
-            return tf.strings.split(inputs).to_tensor()
+            return tf.strings.split(inputs)
 
         def detokenize(self, inputs):
-            return tf.strings.reduce_join([inputs], separator=" ", axis=-1)
+            return tf.strings.reduce_join(inputs, separator=" ", axis=-1)
 
-    tokenizer = WhitespaceTokenizer()
+    tokenizer = WhitespaceSplitterTokenizer()
+
+    # Tokenize some inputs.
+    tokenizer.tokenize("This is a test")
+
+    # Shorthard for `tokenize()`.
+    tokenizer("This is a test")
+
+    # Detokenize some outputs.
+    tokenizer.detokenize(["This", "is", "a", "test"])
     ```
     """
 
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer.py b/keras_nlp/tokenizers/unicode_character_tokenizer.py
index ade085e3..cad5c899 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer.py
@@ -27,6 +27,24 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
     This tokenizer is a vocabulary free tokenizer which tokenizes text as
     unicode characters codepoints.
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
+    If input is a batch of strings (rank > 0):
+    By default, the layer will output a `tf.RaggedTensor` where the last
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`.
+
+    If input is a scalar string (rank == 0):
+    By default, the layer will output a dense `tf.Tensor` with static shape
+    `[None]`. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`.
+
+    The output dtype can be controlled via the `dtype` argument, which should be
+    an integer type (tf.int16, tf.int32, etc.).
+
     Args:
         lowercase: If true, the input text will be first lowered before
             tokenization.
@@ -74,7 +92,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325,    0,    0],
         [1705, 1578, 1575, 1576,    0,    0,    0,    0]], dtype=int32)>
 
-    Tokenize first, then batch the dataset.
+    Tokenize, then batch for ragged outputs.
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
@@ -85,7 +103,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325],
         [1705, 1578, 1575, 1576]]>
 
-    Batch the inputs and then tokenize.
+    Batch, then tokenize for ragged outputs.
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
@@ -95,7 +113,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325],
         [1705, 1578, 1575, 1576]]>
 
-    Tokenize first, then batch for dense outputs (`sequence_length` provided).
+    Tokenize, then batch for dense outputs (`sequence_length` provided).
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -108,8 +126,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340],
         [1705, 1578, 1575, 1576,    0]], dtype=int32)>
 
-    Batch first, then tokenize for dense outputs (`sequence_length` provided).
-    (`sequence_length` provided).
+    Batch, then tokenize for dense outputs (`sequence_length` provided).
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -121,7 +138,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340],
         [1705, 1578, 1575, 1576,    0]], dtype=int32)>
 
-    Tokenization showcasing truncation of long sequences.
+    Tokenization with truncation.
     >>> inputs = ["I Like to Travel a Lot", "मैं किताबें पढ़ना पसंद करता हूं"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -136,7 +153,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
     >>> tokenizer.detokenize(inputs)
     <tf.Tensor: shape=(), dtype=string, numpy=b'ninja'>
 
-    Detokenization while showcasing padded characters being removed
+    Detokenization with padding.
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=7)
     >>> dataset = tf.data.Dataset.from_tensor_slices(["a b c", "b c", "a"])
@@ -150,13 +167,11 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
     <tf.Tensor: shape=(), dtype=string, numpy=b'a b c'>
 
     Detokenization with invalid bytes.
-    >>> # The 10000000 in the inputs tensor below is an invalid value
-    >>> # Hence it replaces to the replacement_char 75 which represents 'K'
     >>> inputs = tf.constant([110, 105, 10000000, 110, 106,  97])
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
-    ...     errors="replace", replacement_char=75)
+    ...     errors="replace", replacement_char=88)
     >>> tokenizer.detokenize(inputs).numpy().decode('utf-8')
-    'niKnja'
+    'niXnja'
     """
 
     def __init__(
@@ -175,9 +190,9 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
             kwargs["dtype"] = tf.int32
         else:
             dtype = tf.dtypes.as_dtype(kwargs["dtype"])
-            if not dtype.is_integer and dtype != tf.string:
+            if not dtype.is_integer:
                 raise ValueError(
-                    "Output dtype must be an integer type of a string. "
+                    "Output dtype must be an integer type. "
                     f"Received: dtype={dtype}"
                 )
 
@@ -251,6 +266,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
             replacement_char=self.replacement_char,
             input_encoding=self.input_encoding,
         )
+        tokens = tf.cast(tokens, self.compute_dtype)
 
         if self.sequence_length:
             output_shape = tokens.shape.as_list()
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index cc1766ec..79ec5057 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -83,12 +83,23 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     `lowercase=False`, `strip_accents=False` and `split=False`. In
     this case, inputs should be pre-split string tensors or ragged tensors.
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
+    If input is a batch of strings (rank > 0):
     By default, the layer will output a `tf.RaggedTensor` where the last
-    dimension of the output is ragged after whitespace splitting and sub-word
-    tokenizing. If `sequence_length` is set, the layer will output a dense
-    `tf.Tensor` where all inputs have been padded or truncated to
-    `sequence_length`. The output dtype can be controlled via the `dtype`
-    argument, which should be either an integer or string type.
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`.
+
+    If input is a scalar string (rank == 0):
+    By default, the layer will output a dense `tf.Tensor` with static shape
+    `[None]`. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`.
+
+    The output dtype can be controlled via the `dtype` argument, which should
+    be either an integer or string type.
 
     Args:
         vocabulary: A list of strings or a string string filename path. If
@@ -115,8 +126,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             must be included in the vocab.
 
     References:
-        - [Schuster and Nakajima, 2012](https://research.google/pubs/pub37842/)
-        - [Song et al., 2020](https://arxiv.org/abs/2012.15524)
+     - [Schuster and Nakajima, 2012](https://research.google/pubs/pub37842/)
+     - [Song et al., 2020](https://arxiv.org/abs/2012.15524)
 
     Examples:
 
diff --git a/keras_nlp/utils/__init__.py b/keras_nlp/utils/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/keras_nlp/utils/__init__.py
@@ -0,0 +1,13 @@
+# 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.
diff --git a/keras_nlp/utils/tensor_utils.py b/keras_nlp/utils/tensor_utils.py
new file mode 100644
index 00000000..26fc815f
--- /dev/null
+++ b/keras_nlp/utils/tensor_utils.py
@@ -0,0 +1,47 @@
+# 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.
+
+import tensorflow as tf
+
+
+def _decode_strings_to_utf8(inputs):
+    """Recursively decodes to list of strings with 'utf-8' encoding."""
+    if isinstance(inputs, bytes):
+        # Handles the case when the input is a scalar string.
+        return inputs.decode("utf-8")
+    else:
+        # Recursively iterate when input is a list.
+        return [_decode_strings_to_utf8(x) for x in inputs]
+
+
+def tensor_to_string_list(inputs):
+    """Detokenize and convert tensor to nested lists of python strings.
+
+    This is a convenience method which converts each byte string to a python
+    string.
+
+    Args:
+        inputs: Input tensor, or dict/list/tuple of input tensors.
+        *args: Additional positional arguments.
+        **kwargs: Additional keyword arguments.
+    """
+    if not isinstance(inputs, (tf.RaggedTensor, tf.Tensor)):
+        inputs = tf.convert_to_tensor(inputs)
+    if isinstance(inputs, tf.RaggedTensor):
+        list_outputs = inputs.to_list()
+    elif isinstance(inputs, tf.Tensor):
+        list_outputs = inputs.numpy()
+        if inputs.shape.rank != 0:
+            list_outputs = list_outputs.tolist()
+    return _decode_strings_to_utf8(list_outputs)
diff --git a/setup.py b/setup.py
index bb458e4c..0b83bca1 100644
--- a/setup.py
+++ b/setup.py
@@ -53,6 +53,7 @@ setup(
             "datasets",  # For GLUE in BERT example.
             "nltk",
             "wikiextractor",
+            "keras-tuner",
         ],
     },
     classifiers=[


Test patch:
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
index cd39f4fc..7d2741be 100644
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -34,7 +34,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length
+            sequence_length=sequence_length
         )
         input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -51,7 +51,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length
+            sequence_length=sequence_length
         )
         input_tensor = tf.keras.Input(
             shape=(feature_size, sequence_length, feature_size)
@@ -75,7 +75,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length, dtype="float16"
+            sequence_length=sequence_length, dtype="float16"
         )
         input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -91,7 +91,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 40
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, feature_size))
@@ -107,7 +107,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 60
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 4-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, None, feature_size))
@@ -123,7 +123,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 40
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, feature_size))
@@ -148,7 +148,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 3
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
@@ -172,7 +172,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 2
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         # Create a 3-dimensional ragged input (the first dimension is implicit).
@@ -209,7 +209,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 2
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         # Create a 4-dimensional ragged input (the first dimension is implicit).
@@ -254,7 +254,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 3
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
@@ -283,11 +283,11 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
     def test_get_config_and_from_config(self):
         max_sequence_length = 40
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         config = test_layer.get_config()
         expected_config_subset = {
-            "max_length": max_sequence_length,
+            "sequence_length": max_sequence_length,
             "initializer": {
                 "class_name": "GlorotUniform",
                 "config": {"seed": None},
@@ -306,7 +306,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 6
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
index a9916ab9..cfaadef6 100644
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -32,7 +32,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
     def test_get_config_and_from_config(self):
         token_and_position_embed = TokenAndPositionEmbedding(
             vocabulary_size=5,
-            max_length=10,
+            sequence_length=10,
             embedding_dim=32,
         )
 
@@ -62,7 +62,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
             embeddings_initializer=custom_embed_init,
         )
@@ -106,7 +106,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
             embeddings_initializer=custom_embed_init,
         )
@@ -124,7 +124,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
     def test_mask_propagation(self):
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=5,
-            max_length=4,
+            sequence_length=4,
             embedding_dim=3,
             mask_zero=True,
         )
@@ -139,7 +139,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
         )
         inputs = keras.Input(shape=(sequence_length,))
diff --git a/keras_nlp/utils/tensor_utils_test.py b/keras_nlp/utils/tensor_utils_test.py
new file mode 100644
index 00000000..d9941f75
--- /dev/null
+++ b/keras_nlp/utils/tensor_utils_test.py
@@ -0,0 +1,34 @@
+# 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.
+
+import tensorflow as tf
+
+from keras_nlp.utils.tensor_utils import tensor_to_string_list
+
+
+class TensorToStringListTest(tf.test.TestCase):
+    def test_detokenize_to_strings_for_ragged(self):
+        input_data = tf.ragged.constant([["▀▁▂▃", "samurai"]])
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertAllEqual(detokenize_output, [["▀▁▂▃", "samurai"]])
+
+    def test_detokenize_to_strings_for_dense(self):
+        input_data = tf.constant([["▀▁▂▃", "samurai"]])
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertAllEqual(detokenize_output, [["▀▁▂▃", "samurai"]])
+
+    def test_detokenize_to_strings_for_scalar(self):
+        input_data = tf.constant("▀▁▂▃")
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertEqual(detokenize_output, "▀▁▂▃")