nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: keras-team/keras-nlp
Pull request #1260 (2023-10-01 12:02:45)

Fix/allow different encoder and decoder feature dimensions in transformer decoder layer
Fixes https://github.com/keras-team/keras-nlp/issues/1235

Thanks @fdtomasi for the detailed issue report, made fixing this quite easy.

mattdangerw: /gcbrun

Patch:
diff --git a/keras_nlp/layers/modeling/transformer_decoder.py b/keras_nlp/layers/modeling/transformer_decoder.py
index 92734fca..3a3cda3f 100644
--- a/keras_nlp/layers/modeling/transformer_decoder.py
+++ b/keras_nlp/layers/modeling/transformer_decoder.py
@@ -199,12 +199,12 @@ class TransformerDecoder(keras.layers.Layer):
             )
             if hasattr(self._cross_attention_layer, "_build_from_signature"):
                 self._cross_attention_layer._build_from_signature(
-                    query=encoder_sequence_shape,
+                    query=decoder_sequence_shape,
                     value=encoder_sequence_shape,
                 )
             else:
                 self._cross_attention_layer.build(
-                    query_shape=encoder_sequence_shape,
+                    query_shape=decoder_sequence_shape,
                     value_shape=encoder_sequence_shape,
                 )
             self._cross_attention_layer_norm = keras.layers.LayerNormalization(
@@ -212,7 +212,7 @@ class TransformerDecoder(keras.layers.Layer):
                 dtype=self.dtype_policy,
                 name="cross_attention_layer_norm",
             )
-            self._cross_attention_layer_norm.build(encoder_sequence_shape)
+            self._cross_attention_layer_norm.build(decoder_sequence_shape)
             self._cross_attention_dropout = keras.layers.Dropout(
                 rate=self.dropout,
                 dtype=self.dtype_policy,


Test patch:
diff --git a/keras_nlp/layers/modeling/transformer_decoder_test.py b/keras_nlp/layers/modeling/transformer_decoder_test.py
index 12cd189e..9a7bc9c1 100644
--- a/keras_nlp/layers/modeling/transformer_decoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_decoder_test.py
@@ -168,3 +168,12 @@ class TransformerDecoderTest(TestCase):
         output, output_cache = call(outputs, input_cache)
         self.assertAllClose(output, no_loop_outputs)
         self.assertAllClose(output_cache, no_loop_cache)
+
+    def test_different_feature_dimension_for_encoder_and_decoder_sequence(self):
+        decoder = TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder_sequence = ops.random.uniform(shape=[1, 4, 6])
+        encoder_sequence = ops.random.uniform(shape=[1, 4, 5])
+        decoder(decoder_sequence, encoder_sequence)