Repository: keras-team/keras-nlp

Issue:
Add a token and position embedding layer
For simplicity of our examples, we should add a layer that combines `keras.layers.Embedding` and `keras_nlp.layers.PositionEmbedding` into a single offering. Initial take for an API signature...

```python
keras_nlp.layers.TokenAndPositionEmbedding(
    vocabulary_size,
    max_length,
    embedding_dim,
    embeddings_initializer="glorot_uniform",
    embeddings_regularizer=None,
    mask_zero=False,
    **kwargs
)
```

Internally, this will combine the layers as follows...

```python
# in __init__
self.token_embeddings = layers.Embedding(
    input_dim=vocab_size, output_dim=embedding_dim
)
self.position_embeddings = keras_nlp.layers.PositionEmbedding(
    max_length=max_length
)

# in call
embedded_tokens = self.token_embeddings(inputs)
embedded_positions = self.position_embeddings(embedded_tokens)
outputs = embedded_tokens + embedded_positions
```


Discussion:
Note that we will do this only for the learned positional embedding (not the sinusoidal encoding). We believe that the learned embedding will be the most common usage for adding position info, so it's the only one we will provide the extra "glue" layer for.
Can I try to work on this issue? Please assign me.
Please do. Thanks!

Resolving patch:
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -15,5 +15,8 @@
 from keras_nlp.layers.fnet_encoder import FNetEncoder
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
+from keras_nlp.layers.token_and_position_embedding import (
+    TokenAndPositionEmbedding,
+)
 from keras_nlp.layers.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -64,7 +64,9 @@ def __init__(
     ):
         super().__init__(**kwargs)
         if max_length is None:
-            raise ValueError("`max_length` must be an Integer, not `None`.")
+            raise ValueError(
+                "`max_length` must be an Integer, received `None`."
+            )
         self.max_length = int(max_length)
         self.initializer = keras.initializers.get(initializer)
 
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -0,0 +1,114 @@
+# 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.
+
+"""Creates an Embedding Layer and adds Positional Embeddings"""
+
+from tensorflow import keras
+
+import keras_nlp.layers
+
+
+class TokenAndPositionEmbedding(keras.layers.Layer):
+    """A layer which sums a token and position embedding.
+
+    This layer assumes that the last dimension in the input corresponds
+    to the sequence dimension.
+
+    Args:
+        vocabulary_size: The size of the vocabulary (should be no larger
+            than 999)
+        max_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
+        mask_zero: Boolean, whether or not the input value 0 is a special
+            "padding" value that should be masked out.
+            This is useful when using recurrent layers which may take variable
+            length input. If this is True, then all subsequent layers in the
+            model need to support masking or an exception will be raised.
+            If mask_zero` is set to True, as a consequence, index 0 cannot be
+            used in the vocabulary
+            (input_dim should equal size of vocabulary + 1).
+
+    Example:
+    ```python
+    seq_length = 50
+    vocab_size = 5000
+    embed_dim = 128
+    inputs = tf.keras.Input(shape=(seq_length,))
+    embedding_layer = keras_nlp.layers.TokenAndPositionEmbedding(
+        vocabulary_size=vocab_size,
+        max_length=max_length,
+        embedding_dim=embed_dim,
+    )
+    outputs = embedding_layer(inputs)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        max_length,
+        embedding_dim,
+        embeddings_initializer="glorot_uniform",
+        mask_zero=False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if vocabulary_size is None:
+            raise ValueError(
+                "`vocabulary_size` must be an Integer, received `None`."
+            )
+        if max_length is None:
+            raise ValueError(
+                "`max_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.embedding_dim = int(embedding_dim)
+        self.token_embedding = keras.layers.Embedding(
+            vocabulary_size,
+            embedding_dim,
+            embeddings_initializer=embeddings_initializer,
+            mask_zero=mask_zero,
+        )
+        self.position_embedding = keras_nlp.layers.PositionEmbedding(
+            max_length=max_length,
+            initializer=embeddings_initializer,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "max_length": self.max_length,
+                "embedding_dim": self.embedding_dim,
+                "embeddings_initializer": keras.initializers.serialize(
+                    self.token_embedding.embeddings_initializer
+                ),
+                "mask_zero": self.token_embedding.mask_zero,
+            },
+        )
+        return config
+
+    def call(self, inputs):
+        embedded_tokens = self.token_embedding(inputs)
+        embedded_positions = self.position_embedding(embedded_tokens)
+        outputs = embedded_tokens + embedded_positions
+        return outputs


Test patch:
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -0,0 +1,182 @@
+# 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 TokenAndPositionEmbedding"""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import TokenAndPositionEmbedding
+
+
+def custom_embed_init(shape, dtype=None):
+    count = 1
+    for length in shape:
+        count *= length
+    return tf.ones(shape, dtype=dtype)
+
+
+class TokenAndPositionEmbeddingTest(tf.test.TestCase):
+    def test_get_config_and_from_config(self):
+        token_and_position_embed = TokenAndPositionEmbedding(
+            vocabulary_size=5,
+            max_length=10,
+            embedding_dim=32,
+        )
+
+        config = token_and_position_embed.get_config()
+
+        expected_config_subset = {
+            "embeddings_initializer": keras.initializers.serialize(
+                keras.initializers.GlorotUniform()
+            ),
+            "mask_zero": False,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
+
+        restored_token_and_position_embed = (
+            TokenAndPositionEmbedding.from_config(config)
+        )
+
+        self.assertEqual(
+            restored_token_and_position_embed.get_config(),
+            {**config, **expected_config_subset},
+        )
+
+    def test_ragged_tensor(self):
+        vocabulary_size = 5
+        sequence_length = 4
+        embedding_dim = 3
+        test_layer = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            max_length=sequence_length,
+            embedding_dim=embedding_dim,
+            embeddings_initializer=custom_embed_init,
+        )
+        # Create a 2-dimensional ragged input
+        # (the first dimension is implicit).
+        input_tensor = tf.keras.Input(
+            shape=(sequence_length,), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = tf.keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.ragged.constant(
+            [
+                [1.0, 1.0],
+                [],
+                [1.0, 1.0, 1.0, 1.0],
+                [1.0, 1.0, 1.0],
+            ],
+            ragged_rank=1,
+        )
+        expected_output_data = tf.ragged.constant(
+            [
+                [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]],
+                [],
+                [
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                ],
+                [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0], [2.0, 2.0, 2.0]],
+            ],
+            ragged_rank=1,
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
+    def test_dense_tensor(self):
+        vocabulary_size = 5
+        sequence_length = 4
+        embedding_dim = 3
+        test_layer = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            max_length=sequence_length,
+            embedding_dim=embedding_dim,
+            embeddings_initializer=custom_embed_init,
+        )
+        # Create a 2-dimensional input
+        # (the first dimension is implicit).
+        input_tensor = tf.keras.Input(
+            shape=(sequence_length,), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = tf.keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.constant(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [1.0, 1.0, 1.0, 1.0],
+            ],
+        )
+        expected_output_data = tf.constant(
+            [
+                [
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                ],
+                [
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                ],
+                [
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                ],
+                [
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                    [2.0, 2.0, 2.0],
+                ],
+            ],
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
+    def test_save_model(self):
+        vocabulary_size = 5
+        sequence_length = 4
+        embedding_dim = 3
+        test_layer = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            max_length=sequence_length,
+            embedding_dim=embedding_dim,
+        )
+        inputs = keras.Input(shape=(sequence_length,))
+        outputs = test_layer(inputs)
+        model = keras.Model(inputs=inputs, outputs=outputs)
+
+        data = tf.zeros(shape=[2, sequence_length])
+        model(data)
+
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path)
+        loaded_model = keras.models.load_model(path)
+
+        model_output = model.predict(data)
+        loaded_model_output = loaded_model.predict(data)
+        self.assertAllClose(model_output, loaded_model_output)