Repository: keras-team/keras (Python)
Task: Bug Fix (2024-03-11 17:59:37+00:00)

Issue:
Keras 3 Attention layer value tensor dimension
hi,

I found the  below would not return the proper size output in Keras 3 (but works fine in Keras 2)
Please help to fix it, 

Thanks.



```python
import keras
from keras import layers

i = layers.Input((8,4))
xq = layers.Conv1D(5,1)(i)
xk = layers.Conv1D(5,1)(i)
xv = layers.Conv1D(7,1)(i)
o = layers.Attention()([xq,xv,xk])

m = keras.Model(inputs=i, outputs=o)
m.summary()

Output as below
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)        ┃ Output Shape      ┃    Param # ┃ Connected to      ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ input_layer         │ (None, 8, 4)      │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv1d (Conv1D)     │ (None, 8, 5)      │         25 │ input_layer[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv1d_2 (Conv1D)   │ (None, 8, 7)      │         35 │ input_layer[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv1d_1 (Conv1D)   │ (None, 8, 5)      │         25 │ input_layer[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ attention           │ **(None, 8, 5)**      │          0 │ conv1d[0][0],     │
│ (Attention)         │                   │            │ conv1d_2[0][0],   │
│                     │                   │            │ conv1d_1[0][0]    │
└─────────────────────┴───────────────────┴────────────┴───────────────────┘
 Total params: 85 (340.00 B)
 Trainable params: 85 (340.00 B)
 Non-trainable params: 0 (0.00 B)
```

The Attention layer output shape should be (None, 8, 7), since **xv** is from **Conv1D** with 7 kernels.

Thie same code gives correct output from keras 2

```python
Model: "model"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 input_1 (InputLayer)           [(None, 8, 4)]       0           []                               
                                                                                                  
 conv1d (Conv1D)                (None, 8, 5)         25          ['input_1[0][0]']                
                                                                                                  
 conv1d_2 (Conv1D)              (None, 8, 7)         35          ['input_1[0][0]']                
                                                                                                  
 conv1d_1 (Conv1D)              (None, 8, 5)         25          ['input_1[0][0]']                
                                                                                                  
 attention (Attention)          **(None, 8, 7)**         0           ['conv1d[0][0]',                 
                                                                  'conv1d_2[0][0]',               
                                                                  'conv1d_1[0][0]']               
                                                                                                  
==================================================================================================
Total params: 85
Trainable params: 85
Non-trainable params: 0
```



Discussion:

Patch:
diff --git a/keras/dtype_policies/dtype_policy.py b/keras/dtype_policies/dtype_policy.py
--- a/keras/dtype_policies/dtype_policy.py
+++ b/keras/dtype_policies/dtype_policy.py
@@ -173,9 +173,6 @@ def _parse_name(self, name):
             return "float16", "float32"
         elif name == "mixed_bfloat16":
             return "bfloat16", "float32"
-        elif name == "uint8":
-            dtype = backend.standardize_dtype(name)
-            return dtype, dtype
         try:
             dtype = backend.standardize_dtype(name)
             return dtype, dtype
diff --git a/keras/layers/attention/attention.py b/keras/layers/attention/attention.py
--- a/keras/layers/attention/attention.py
+++ b/keras/layers/attention/attention.py
@@ -242,7 +242,8 @@ def compute_mask(self, inputs, mask=None):
         return ops.convert_to_tensor(mask[0])
 
     def compute_output_shape(self, input_shape):
-        return input_shape[0]
+        """Returns shape of value tensor dim, but for query tensor length"""
+        return (*input_shape[0][:-1], input_shape[1][-1])
 
     def _validate_inputs(self, inputs, mask=None):
         """Validates arguments of the call method."""


Test patch:
diff --git a/keras/layers/attention/attention_test.py b/keras/layers/attention/attention_test.py
--- a/keras/layers/attention/attention_test.py
+++ b/keras/layers/attention/attention_test.py
@@ -342,3 +342,19 @@ def test_attention_compute_mask_with_different_input_shapes(self):
             computed_mask = layer.compute_mask(inputs=dummy_inputs, mask=mask)
             computed_mask = ops.convert_to_numpy(computed_mask)
             self.assertTrue(np.array_equal(computed_mask, valid_mask))
+
+    def test_attention_compute_output_shape(self):
+        layer = layers.Attention()
+
+        query = np.random.random((2, 3, 4))
+        value = np.random.random((2, 3, 5))
+        key = np.random.random((2, 3, 4))
+        layer = layers.Attention()
+        output = layer([query, value, key])
+        self.assertAllEqual(output.shape, value.shape)
+        self.assertAllEqual(
+            layer.compute_output_shape(
+                input_shape=[query.shape, value.shape, key.shape]
+            ),
+            output.shape,
+        )