Repository: keras-team/keras-nlp

Issue:
Move RotaryEmbedding to layers
We need to move `RotaryEmbedding` to layers and add unit tests once #1056 is merged


Discussion:

Resolving patch:
diff --git a/keras_nlp/layers/modeling/rotary_embedding.py b/keras_nlp/layers/modeling/rotary_embedding.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/modeling/rotary_embedding.py
@@ -0,0 +1,128 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+
+
+@keras_nlp_export("keras_nlp.layers.RotaryEmbedding")
+class RotaryEmbedding(keras.layers.Layer):
+    """Rotary positional encoding layer.
+
+    This layer encodes absolute positional information with a rotation
+    matrix. It calculates the rotary encoding with a mix of sine and
+    cosine functions with geometrically increasing wavelengths.
+    Defined and formulated in [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4).
+    The input must be a tensor with shape a sequence dimension and a feature
+    dimension. Typically, this will either an input with shape
+    `(batch_size, sequence_length, feature_length)` or
+    `(batch_size, sequence_length, num_heads, feature_length)`.
+    This layer will return a new tensor with the rotary embedding applied to
+    the input tensor.
+
+    Args:
+        max_wavelength: int. The maximum angular wavelength of the sine/cosine
+            curves.
+        scaling_factor: float. The scaling factor used to scale frequency range.
+        sequence_axis: int. Sequence axis in the input tensor.
+        feature_axis: int. Feature axis in the input tensor.
+
+    Examples:
+
+    ```python
+    batch_size = 16
+    feature_length = 18
+    sequence_length = 256
+    num_heads = 8
+
+    # No multi-head dimension.
+    tensor = tf.ones((batch_size, sequence_length, feature_length))
+    rot_emb_layer = RotaryEmbedding()
+    tensor_rot = rot_emb_layer(tensor)
+
+    # With multi-head dimension.
+    tensor = tf.ones((batch_size, sequence_length, num_heads, feature_length))
+    tensor_rot = rot_emb_layer(tensor)
+    ```
+
+    References:
+     - [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4)
+    """
+
+    def __init__(
+        self,
+        max_wavelength=10000,
+        scaling_factor=1.0,
+        sequence_axis=1,
+        feature_axis=-1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.max_wavelength = max_wavelength
+        self.sequence_axis = sequence_axis
+        self.feature_axis = feature_axis
+        self.scaling_factor = scaling_factor
+
+    def call(self, inputs):
+        rotary_dim = ops.shape(inputs)[-1]
+        cos_emb, sin_emb = self._compute_cos_sin_embedding(inputs, rotary_dim)
+        return self._apply_rotary_pos_emb(inputs, cos_emb, sin_emb)
+
+    def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
+        x1, x2 = ops.split(tensor, 2, axis=self.feature_axis)
+        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):
+        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(
+            self.scaling_factor, self.compute_dtype
+        )
+        inverse_freq = 1.0 / (
+            self.max_wavelength
+            ** (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.cast(tensor, dtype=inverse_freq.dtype)
+        freq = ops.einsum("i, j -> ij", tensor, inverse_freq)
+        embedding = ops.concatenate((freq, freq), axis=self.feature_axis)
+
+        def get_axis(axis):
+            return axis if axis > 0 else len(x.shape) + axis
+
+        feature_axis = get_axis(self.feature_axis)
+        sequence_axis = get_axis(self.sequence_axis)
+
+        for axis in range(len(x.shape)):
+            if axis != sequence_axis and axis != feature_axis:
+                embedding = ops.expand_dims(embedding, axis)
+
+        return ops.cos(embedding), ops.sin(embedding)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_wavelength": self.max_wavelength,
+                "scaling_factor": self.scaling_factor,
+                "sequence_axis": self.sequence_axis,
+                "feature_axis": self.feature_axis,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
@@ -14,7 +14,7 @@
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.models.gpt_neo_x.rotary_embedding import RotaryEmbedding
+from keras_nlp.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_nlp.utils.keras_utils import clone_initializer
 
 
@@ -63,7 +63,9 @@ def __init__(
         self.attn_head_size = hidden_dim // num_heads
         self.rotary_max_wavelength = rotary_max_wavelength
         self.rotary_dim = int(self.attn_head_size * rotary_percentage)
-        self.rotary_embedding = RotaryEmbedding(rotary_max_wavelength)
+        self.rotary_embedding_layer = RotaryEmbedding(
+            max_wavelength=rotary_max_wavelength
+        )
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self.max_sequence_length = max_sequence_length
@@ -190,8 +192,8 @@ def call(
             key[..., self.rotary_dim :],
         )
 
-        query_rot = self.rotary_embedding(query_rot)
-        key_rot = self.rotary_embedding(key_rot)
+        query_rot = self.rotary_embedding_layer(query_rot)
+        key_rot = self.rotary_embedding_layer(key_rot)
 
         query = ops.concatenate((query_rot, query_pass), axis=-1)
         key = ops.concatenate((key_rot, key_pass), axis=-1)
diff --git a/keras_nlp/models/gpt_neo_x/rotary_embedding.py b/keras_nlp/models/gpt_neo_x/rotary_embedding.py
deleted file mode 100644
--- a/keras_nlp/models/gpt_neo_x/rotary_embedding.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# Copyright 2023 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from keras_nlp.backend import keras
-from keras_nlp.backend import ops
-
-
-class RotaryEmbedding(keras.layers.Layer):
-    def __init__(self, max_wavelength=10000, **kwargs):
-        super().__init__(**kwargs)
-        self.max_wavelength = max_wavelength
-
-    def call(self, inputs):
-        rotary_dim = ops.shape(inputs)[-1]
-
-        cos_emb, sin_emb = self._compute_cos_sin_embedding(
-            inputs, rotary_dim, seq_dim=1
-        )
-        outputs = self._apply_rotary_pos_emb(inputs, cos_emb, sin_emb)
-
-        return outputs
-
-    def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
-        cos_emb = cos_emb[:, : ops.shape(tensor)[1], :, :]
-        sin_emb = sin_emb[:, : ops.shape(tensor)[1], :, :]
-
-        x1, x2 = ops.split(tensor, 2, axis=-1)
-        half_rot_tensor = ops.concatenate((-x2, x1), axis=-1)
-
-        return (tensor * cos_emb) + (half_rot_tensor * sin_emb)
-
-    def _compute_cos_sin_embedding(self, x, rotary_dim, seq_dim=1):
-        seq_len = ops.shape(x)[seq_dim]
-
-        range = ops.arange(0, rotary_dim, 2, self.compute_dtype)
-        inverse_freq = 1.0 / (
-            self.max_wavelength
-            ** (range / ops.cast(rotary_dim, self.compute_dtype))
-        )
-
-        tensor = ops.arange(seq_len, dtype=inverse_freq.dtype)
-        freqs = ops.einsum("i, j -> ij", tensor, inverse_freq)
-        embedding = ops.concatenate((freqs, freqs), axis=-1)[None, :, None, :]
-
-        return ops.cos(embedding), ops.sin(embedding)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "max_wavelength": self.max_wavelength,
-            }
-        )
-        return config
-
-    def compute_output_shape(self, input_shape):
-        return input_shape


Test patch:
diff --git a/keras_nlp/layers/modeling/rotary_embedding_test.py b/keras_nlp/layers/modeling/rotary_embedding_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/modeling/rotary_embedding_test.py
@@ -0,0 +1,114 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+from keras_nlp.layers.modeling.rotary_embedding import RotaryEmbedding
+from keras_nlp.tests.test_case import TestCase
+
+
+class RotaryEmbeddingTest(TestCase):
+    def test_valid_call(self):
+        embedding_layer = RotaryEmbedding()
+        model = keras.Sequential(
+            [
+                keras.Input(shape=(4, 6)),
+                embedding_layer,
+            ]
+        )
+        input = ops.ones(shape=[2, 4, 6])
+        model(input)
+
+    def test_static_layer_output_shape(self):
+        embedding_layer = RotaryEmbedding()
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(seq_length, hidden_size))
+        outputs = embedding_layer(inputs)
+
+        # When using static positional encoding shapes, the output is expected
+        # to be the same as the input shape in all dimensions.
+        expected_output_shape = (None, seq_length, hidden_size)
+        self.assertEqual(expected_output_shape, outputs.shape)
+
+    def test_dynamic_layer_output_shape(self):
+        embedding_layer = RotaryEmbedding()
+        hidden_size = 32
+        inputs = keras.Input(shape=(None, hidden_size))
+        outputs = embedding_layer(inputs)
+
+        # When using dynamic positional encoding shapes, the output is expected
+        # to be the same as the input shape in all dimensions but may be None.
+        expected_output_shape = (None, None, hidden_size)
+        self.assertEqual(expected_output_shape, outputs.shape)
+
+    # do multi dimension before sequence length
+    def test_multi_dimension_layer_output_shape(self):
+        embedding_layer = RotaryEmbedding()
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(None, seq_length, hidden_size))
+        outputs = embedding_layer(inputs)
+
+        # When using muliple dimensions before sequence length, the output is
+        # expected to be the same as the input shape in all dimensions.
+        expected_output_shape = (None, None, seq_length, hidden_size)
+        self.assertEqual(expected_output_shape, outputs.shape)
+
+    def test_output_correct_values(self):
+        embedding_layer = RotaryEmbedding()
+        model = keras.Sequential(
+            [
+                keras.Input(shape=(4, 6)),
+                embedding_layer,
+            ]
+        )
+        input = ops.ones(shape=[1, 4, 6])
+        output = model(input)
+
+        # comapre position encoding values for position 0 and 3
+        expected_0 = [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]
+        expected_3 = [-1.1311, 0.8515, 0.9935, -0.8489, 1.1291, 1.0064]
+        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_get_config_and_from_config(self):
+        embedding_layer = RotaryEmbedding(
+            max_wavelength=1000,
+            scaling_factor=2.0,
+            sequence_axis=1,
+            feature_axis=-1,
+        )
+        config = embedding_layer.get_config()
+        expected_config = {
+            "max_wavelength": 1000,
+            "scaling_factor": 2.0,
+            "sequence_axis": 1,
+            "feature_axis": -1,
+        }
+        self.assertEqual(config, {**config, **expected_config})
+        restored_embedding_layer = RotaryEmbedding.from_config(config)
+        self.assertEqual(
+            restored_embedding_layer.get_config(),
+            {**config, **expected_config},
+        )
+
+    def test_float16_dtype(self):
+        embedding_layer = RotaryEmbedding(dtype="float16")
+        seq_length = 100
+        hidden_size = 32
+        inputs = keras.Input(shape=(seq_length, hidden_size))
+        outputs = embedding_layer(inputs)
+
+        # output dtype for this layer should be float16.
+        self.assertEqual(outputs.dtype, "float16")