Repository: keras-team/keras

Issue:
Implement `compute_output_shape` for `BaseAttentionLayer`

**System information**
- TensorFlow version (you are using): 2.8
- Are you willing to contribute it (Yes/No): yes


**Describe the feature and the current behavior/state.**
Currently the method compute_output_shape is not implemented for layers derived from `BaseAttentionLayer`. Hence, it's impossible to use these layers with, for example, `TimeDistributed` wrapper (see [issue](https://github.com/keras-team/keras/issues/15515) Attention module not working with TimeDistributed layer).
To fix this, one could use eager mode or implement some workaround. Another way is to implement `compute_output_shape` method for `BaseAttentionLayer`. But `BaseAttentionLayer.call` method has a parameter `return_attention_weights` that adds attention weights tensor to values returned by attention layer, if set. It could change output shape of the layer after build.
Moreover, parameter `return_attention_weights` does not change a state of the layer when provided and is not saved to the object's attributes. So, it's impossible to implement `compute_output_shape` method while `return_attention_weights` belongs to the parameters of the call method. 

My suggestion is to change the API like this:

from 
```python3
class BaseDenseAttention(Layer):
    def __init__(self, causal=False, dropout=0.0,
                 **kwargs):
      ...
    
    def call(self,
             inputs,
             mask=None,
             training=None,
             return_attention_scores=False):
      ...
```

to

```python3
class BaseDenseAttention(Layer):
    def __init__(self, causal=False, dropout=0.0, return_attention_scores=False),  # move the parameter to the constructor
                 **kwargs):
      self.return_attention_scores=return_attention_scores
      ...
    
    def call(self, 
             inputs,
             mask=None,
             training=None,
             return_attention_scores=False):   # <- make it deprecated and add deprecation warning, remove in future versions
      ...

  def compute_output_shape(self, input_shape):
      # this function returns output shape using self.return_attention_scores as a conditional 
      ...

These changes make computation of output shape possible, but keep the possibility to return attention weights, if it's necessary for model debugging and/or analysis. For most use cases, it's enough to place `return_attention_scores` in a constructor of a class. Likewise, in other keras layers the parameters that change output shape of the layer are placed in a constructor of the layer (for example, `return_sequences` in recurrent layers). 
```

Related issues:
- [Unable to create](https://github.com/keras-team/keras/issues/15515) TimeDistributed wrapper for Attention layers
- [Feature request](https://github.com/tensorflow/tensorflow/issues/44127) for the parameter `return_attention_weights` to the `call` method. 

**Will this change the current api? How?**
Parameter return_attention_weights will be moved from call method of BaseAttentionLayer layers to constructor of BaseAttentionLayer. That allows to compute output shape of the attention layers and use them with, for example, TimeDistributed wrapper (that fixes previously mentioned issue).

**Who will benefit with this feature?**
Anyone who would use attention layers.

**Any Other info.**
It could break some of the tutorials/examples. Fixes may be required in the future: [this NMT tutorial](https://github.com/tensorflow/text/blob/master/docs/tutorials/nmt_with_attention.ipynb) uses `return_attention_scores` parameter.


Discussion:

Resolving patch:
diff --git a/keras/layers/dense_attention.py b/keras/layers/dense_attention.py
--- a/keras/layers/dense_attention.py
+++ b/keras/layers/dense_attention.py
@@ -58,7 +58,7 @@ class BaseDenseAttention(base_layer.BaseRandomLayer):
         `mask==False` do not contribute to the result.
     training: Python boolean indicating whether the layer should behave in
       training mode (adding dropout) or in inference mode (no dropout).
-    return_attention_scores: bool, it `True`, returns the attention scores
+    return_attention_scores: bool, if `True`, returns the attention scores
       (after masking and softmax) as an additional output argument.
 
   Output:
@@ -180,6 +180,11 @@ def compute_mask(self, inputs, mask=None):
       return tf.convert_to_tensor(q_mask)
     return None
 
+  def compute_output_shape(self, input_shape):
+    # return_attention_scores argument of BaseDenseAttention.call method
+    # is ignored. Output shape of attention_scores cannot be returned.
+    return tf.TensorShape(input_shape[0])
+
   def _validate_call_args(self, inputs, mask):
     """Validates arguments of the call method."""
     class_name = self.__class__.__name__


Test patch:
diff --git a/keras/layers/wrappers_test.py b/keras/layers/wrappers_test.py
--- a/keras/layers/wrappers_test.py
+++ b/keras/layers/wrappers_test.py
@@ -514,6 +514,28 @@ def compute_output_shape(self, input_shape):
         epochs=1,
         batch_size=3)
 
+  def test_TimeDistributed_Attention(self):
+     query_input = keras.layers.Input(shape=(None, 1, 10), dtype='float32')
+     value_input = keras.layers.Input(shape=(None, 4, 10), dtype='float32')
+
+     # Query-value attention of shape [batch_size, Tq, filters].
+     query_value_attention_seq = keras.layers.TimeDistributed(
+       keras.layers.Attention())([query_input, value_input])
+     model = keras.models.Model(
+       [query_input, value_input], query_value_attention_seq)
+     model.compile(optimizer='rmsprop', loss='mse')
+     model.fit(
+       [np.random.random((10, 8, 1, 10)), np.random.random((10, 8, 4, 10))],
+       np.random.random((10, 8, 1, 10)),
+       epochs=1,
+       batch_size=10
+     )
+
+     # test config and serialization/deserialization
+     model.get_config()
+     model = keras.models.model_from_json(model.to_json())
+     model.summary()
+
 
 @combinations.generate(combinations.combine(mode=['graph', 'eager']))
 class BidirectionalTest(tf.test.TestCase, parameterized.TestCase):