Add a test against keras-nightly

.github/workflows/actions.yml

@@ -17,10 +17,12 @@ jobs:
       fail-fast: false
       matrix:
         backend: [tensorflow, jax, torch]
-        version: [latest]
+        version: [stable]
         include:
-          - backend: torch
+          - backend: jax
             version: 3.1
+          - backend: jax
+            version: nightly
     runs-on: ubuntu-latest
     env:
       KERAS_BACKEND: ${{ matrix.backend }}
@@ -46,11 +48,16 @@ jobs:
       run: |
           pip install -r requirements.txt --progress-bar off
           pip install --no-deps -e "." --progress-bar off
-    - name: Pin Keras version
+    - name: Pin Keras 3.1
       if: ${{ matrix.version == '3.1'}}
       run: |
         pip uninstall -y keras
         pip install keras==3.1.0 --progress-bar off
+    - name: Pin Keras Nightly
+      if: ${{ matrix.version == 'nightly'}}
+      run: |
+        pip uninstall -y keras
+        pip install keras-nightly --progress-bar off
     - name: Test with pytest
       run: |
         pytest keras_nlp/

keras_nlp/src/layers/modeling/transformer_decoder_test.py

@@ -237,12 +237,3 @@ def call(outputs, self_cache, cross_cache):
         self.assertAllClose(output, no_loop_outputs)
         self.assertAllClose(self_cache, no_loop_self_cache)
         self.assertAllClose(cross_cache, no_loop_cross_cache)
-
-    def test_different_feature_dimension_for_encoder_and_decoder_sequence(self):
-        decoder = TransformerDecoder(
-            intermediate_dim=4,
-            num_heads=2,
-        )
-        decoder_sequence = random.uniform(shape=[1, 4, 6])
-        encoder_sequence = random.uniform(shape=[1, 4, 5])
-        decoder(decoder_sequence, encoder_sequence)