Rename max_length to sequence_length

examples/bert/bert_model.py

@@ -484,7 +484,7 @@ def __init__(
 
         self._position_embedding_layer = keras_nlp.layers.PositionEmbedding(
             initializer=initializer,
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             name="position_embedding",
         )
 

keras_nlp/layers/position_embedding.py

@@ -32,7 +32,7 @@ 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".
         seq_axis: The axis of the input tensor where we add the embeddings.
@@ -43,7 +43,7 @@ class PositionEmbedding(keras.layers.Layer):
         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)
@@ -58,23 +58,23 @@ class PositionEmbedding(keras.layers.Layer):
 
     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 +84,7 @@ def build(self, input_shape):
         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,
         )

keras_nlp/layers/position_embedding_test.py

@@ -34,7 +34,7 @@ def test_static_layer_output_shape(self):
         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 @@ def test_more_than_3_dimensions_static(self):
         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 @@ def test_float16_dtype(self):
         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 @@ def test_dynamic_layer_output_shape(self):
         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 @@ def test_more_than_3_dimensions_dynamic(self):
         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 @@ def test_dynamic_layer_slicing(self):
         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 @@ def test_callable_initializer(self):
         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 @@ def test_ragged_tensor_with_3_dimensions(self):
         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 @@ def test_ragged_tensor_with_4_dimensions(self):
         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 @@ def test_one_training_step(self):
         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 @@ def test_one_training_step(self):
     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 @@ def test_save_model(self):
         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)

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
@@ -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 @@ def __init__(
             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 @@ def __init__(
             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 @@ def get_config(self):
         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

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 @@ def test_ragged_tensor(self):
         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 @@ def test_dense_tensor(self):
         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 @@ def test_dense_tensor(self):
     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 @@ def test_save_model(self):
         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,))