Repository: keras-team/keras-nlp
Pull request #1209 (2023-08-15 16:59:59)

Add start_index argument to all position embedding layers
This will allow all the layers to be used in a decoding context.

mattdangerw: /gcbrun
mattdangerw: /gcbrun

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


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))