fix attention output with symbolic tensors and attention scores (#20689)

Author: Surya2k1

keras/src/layers/attention/attention.py

@@ -280,7 +280,7 @@ def compute_output_spec(
         output_spec = KerasTensor(output_shape, dtype=self.compute_dtype)
 
         # Handle attention scores if requested
-        if self._return_attention_scores:
+        if self._return_attention_scores or return_attention_scores:
             scores_shape = (
                 query.shape[0],
                 query.shape[1],

keras/src/layers/attention/attention_test.py

@@ -417,3 +417,15 @@ def test_return_attention_scores_true_tuple_then_unpack(self):
         self.assertEqual(
             attention_scores.shape, (2, 8, 4)
         )  # Attention scores shape
+
+    def test_return_attention_scores_with_symbolic_tensors(self):
+        """Test to check outputs with symbolic tensors with
+        return_attention_scores = True"""
+        attention = layers.Attention()
+        x = layers.Input(shape=(3, 5))
+        y = layers.Input(shape=(4, 5))
+        output, attention_scores = attention(
+            [x, y], return_attention_scores=True
+        )
+        self.assertEqual(output.shape, (None, 3, 5))  # Output shape
+        self.assertEqual(attention_scores.shape, (None, 3, 4))