Repository: tensorflow/addons

Issue:
Error in types description of tfa.seq2seq.AttentionWrapper
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.3 LT
- TensorFlow version and how it was installed (source or binary): 2.2.0
- TensorFlow-Addons version and how it was installed (source or binary): 0.10.0
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

As I noticed you added type description to classes.
But I found that at least in my case there is a conflict of types in description with real mission of arguments.
According to docs _attention_mechanism_ and _attention_layer_size_ in **tfa.seq2seq.AttentionWrapper** could be lists (for MultiHead attention case) or single instances. But in types you specified that _attention_mechanism_ and _attention_layer_size_ must be **tf.keras.layers.Layer** and  **Optional[FloatTensorLike]** respectively.

https://github.com/tensorflow/addons/blob/v0.9.1/tensorflow_addons/seq2seq/attention_wrapper.py#L1621



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/attention_wrapper.py b/tensorflow_addons/seq2seq/attention_wrapper.py
--- a/tensorflow_addons/seq2seq/attention_wrapper.py
+++ b/tensorflow_addons/seq2seq/attention_wrapper.py
@@ -28,10 +28,11 @@
     FloatTensorLike,
     TensorLike,
     Initializer,
+    Number,
 )
 
 from typeguard import typechecked
-from typing import Optional, Callable
+from typing import Optional, Callable, Union, List
 
 # TODO: Find public API alternatives to these
 from tensorflow.python.keras.engine import base_layer_utils
@@ -1618,14 +1619,16 @@ class AttentionWrapper(tf.keras.layers.AbstractRNNCell):
     def __init__(
         self,
         cell: tf.keras.layers.Layer,
-        attention_mechanism: tf.keras.layers.Layer,
-        attention_layer_size: Optional[FloatTensorLike] = None,
+        attention_mechanism: Union[tf.keras.layers.Layer, List[tf.keras.layers.Layer]],
+        attention_layer_size: Optional[Union[Number, List[Number]]] = None,
         alignment_history: bool = False,
         cell_input_fn: Optional[Callable] = None,
         output_attention: bool = True,
         initial_cell_state: Optional[TensorLike] = None,
         name: Optional[str] = None,
-        attention_layer: Optional[tf.keras.layers.Layer] = None,
+        attention_layer: Optional[
+            Union[tf.keras.layers.Layer, List[tf.keras.layers.Layer]]
+        ] = None,
         attention_fn: Optional[Callable] = None,
         **kwargs
     ):
@@ -1701,7 +1704,7 @@ def __init__(
             and cell output as inputs to generate attention at each time step.
             If None (default), use the context as attention at each time step.
             If attention_mechanism is a list, attention_layer must be a list of
-            the same length. If attention_layers_size is set, this must be
+            the same length. If attention_layer_size is set, this must be
             None.
           attention_fn: An optional callable function that allows users to
             provide their own customized attention function, which takes input


Test patch:
diff --git a/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py b/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
--- a/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
+++ b/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
@@ -932,3 +932,15 @@ def test_attention_wrapper_with_gru_cell():
     initial_state = cell.get_initial_state(inputs=inputs)
     _, state = cell(inputs, initial_state)
     tf.nest.assert_same_structure(initial_state, state)
+
+
+def test_attention_wrapper_with_multiple_attention_mechanisms():
+    cell = tf.keras.layers.LSTMCell(5)
+    mechanisms = [wrapper.LuongAttention(units=3), wrapper.LuongAttention(units=3)]
+    # We simply test that the wrapper creation makes no error.
+    wrapper.AttentionWrapper(cell, mechanisms, attention_layer_size=[4, 5])
+    wrapper.AttentionWrapper(
+        cell,
+        mechanisms,
+        attention_layer=[tf.keras.layers.Dense(4), tf.keras.layers.Dense(5)],
+    )