Repository: tensorflow/addons

Issue:
AttentionWrapper with non LSTM cells returns ValueError
**System information**
- OS Platform and Distribution: Ubuntu 20.04
- 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.7.7
- Is GPU used? (yes/no): yes

**Describe the bug**

Working with @jimthompson5802 on this PR https://github.com/uber/ludwig/pull/699 we discovered an issue with seq2seq attention. 
In seq2seq decoding with attention, one has to wrap the recurrent cell with an AttentionWrapper.
Right now if you do so with any cell type other than LSTMCell (so with SimpleRNNCell and GRUCell), you get a:

```
ValueError: The two structures don't have the same nested structure.

First structure: type=AttentionWrapperState str=AttentionWrapperState(cell_state=<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>, attention=<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>, alignments=<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>, alignment_history=(), attention_state=<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>)

Second structure: type=AttentionWrapperState str=AttentionWrapperState(cell_state=[<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>], attention=<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>, alignments=<tf.Tensor: shape=(128, 128), dtype=float32, dtype=float32)>, alignment_history=(), attention_state=<tf.Tensor: shape=(128, 128), dtype=float32, dtype=float32)>)

More specifically: Substructure "type=list str=[<tf.Tensor: shape=(128, 256), dtype=float32, dtype=float32)>]" is a sequence, while substructure "type=EagerTensor str=tf.Tensor(
[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]], shape=(128, 256), dtype=float32)" is not
Entire first structure:
AttentionWrapperState(cell_state=., attention=., alignments=., alignment_history=(), attention_state=.)
Entire second structure:
AttentionWrapperState(cell_state=[.], attention=., alignments=., alignment_history=(), attention_state=.)
```
This suggests that the AttentionWrapper works well with LSTMCell because it returns a state that's a list of two tensors, but it doesn't work with SimpleRNN and GRU because the state they return is just a Tensor, not a list of length 1 containing one tensor.

I believe there are 2 solutions: one is making all cells return states as a list of tensors. This will be good and it will guarantee homogeneity. The other option is to correct the AttentionWrapper to work both with cells that retunr lists of tensors as states like LSTM and cells that return just a tensor as state like SimpleRNN and GRU.

For more context about the use of this in the original PR, plese refer at the comments starting from  [this one](https://github.com/uber/ludwig/pull/699#issuecomment-629736171).

**Code to reproduce the issue**

Here is a script to reproduce the issue.

```
import numpy as np
import tensorflow as tf
from tensorflow import keras
import tensorflow_addons as tfa

# ValueError exception raised for both SimpleRNNCell and GRUCell

RNN_CELL_CLASS = keras.layers.SimpleRNNCell  # Raises ValueError Exception
# RNN_CELL_CLASS = keras.layers.GRUCell      # Raises ValueError Exception
# RNN_CELL_CLASS = keras.layers.LSTMCell     # Does not raise an error

VOCAB_SIZE = 100
EMBED_SIZE = 10
RNN_UNITS = 256
INPUT_SEQUENCE_SIZE = 10
OUTPUT_SEQUENCE_SIZE = 15
BATCH_SIZE = 128

targets = tf.convert_to_tensor(
    np.random.randint(VOCAB_SIZE, size=OUTPUT_SEQUENCE_SIZE * BATCH_SIZE)
    .reshape(BATCH_SIZE, OUTPUT_SEQUENCE_SIZE),
    dtype=tf.float32
)
encoder_outputs = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
if RNN_CELL_CLASS._keras_api_names[0] == 'keras.layers.LSTMCell':
    encoder_state = [tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32),
                     tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)]
else:
    encoder_state = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
    # note this returns error also if:
    # encoder_state = [tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)]

embeddings_dec = keras.layers.Embedding(VOCAB_SIZE, EMBED_SIZE)
decoder_cell = RNN_CELL_CLASS(RNN_UNITS)
output_layer = keras.layers.Dense(VOCAB_SIZE)

sampler = tfa.seq2seq.sampler.TrainingSampler()
attention_mechanism = tfa.seq2seq.LuongAttention(units=RNN_UNITS)
#attention_mechanism = tfa.seq2seq.BahdanauAttention(units=RNN_UNITS)  # gives same error
decoder_cell = tfa.seq2seq.AttentionWrapper(
    decoder_cell,
    attention_mechanism,
    attention_layer_size=RNN_UNITS
)

decoder = tfa.seq2seq.basic_decoder.BasicDecoder(
    decoder_cell,
    sampler,
    output_layer=output_layer
)

# setup for Attention
attention_mechanism.setup_memory(
    encoder_outputs,
    memory_sequence_length=tf.ones(
        [BATCH_SIZE], dtype=tf.int32) * INPUT_SEQUENCE_SIZE
)

decoder_initial_state = decoder_cell.get_initial_state(
    batch_size=BATCH_SIZE,
    dtype=tf.float32
)
decoder_initial_state = decoder_initial_state.clone(
    cell_state=encoder_state
)

decoder_emb_inp = embeddings_dec(targets)

# Following statement generates exception
# ValueError: The two structures don't have the same nested structure
# for SimpleRNNCell and GRUCell
(
    final_outputs,
    final_state,
    final_sequence_lengths
) = decoder(
    decoder_emb_inp,
    initial_state=decoder_initial_state,
    sequence_length=tf.ones([BATCH_SIZE],
                            dtype=tf.int32) * OUTPUT_SEQUENCE_SIZE
)

```

Note that changing the value of `RNN_CELL` changes the outcome:

```
RNN_CELL_CLASS = keras.layers.SimpleRNNCell  # Raises ValueError Exception
# RNN_CELL_CLASS = keras.layers.GRUCell      # Raises ValueError Exception
# RNN_CELL_CLASS = keras.layers.LSTMCell     # Does not raise an error
```

Moreover, assigning `encoder_state both as a tensor or as a list does not work:

```
encoder_state = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
# encoder_state = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
```

Finally, also note that this is independnent of the type of attention used, the same error is return for both LuongAttention and BahdanauAttention.

Also here's a version of the same script where attention is not used, it works for all different kinds of cells:

```
import numpy as np
import tensorflow as tf
from tensorflow import keras
import tensorflow_addons as tfa

RNN_CELL_CLASS = keras.layers.SimpleRNNCell    # Does not raise an error
# RNN_CELL_CLASS = keras.layers.GRUCell        # Does not raise an error
# RNN_CELL_CLASS = keras.layers.LSTMCell       # Does not raise an error

VOCAB_SIZE = 100
EMBED_SIZE = 10
RNN_UNITS = 256
INPUT_SEQUENCE_SIZE = 10
OUTPUT_SEQUENCE_SIZE = 15
BATCH_SIZE = 128

targets = tf.convert_to_tensor(
    np.random.randint(VOCAB_SIZE, size=OUTPUT_SEQUENCE_SIZE * BATCH_SIZE)
    .reshape(BATCH_SIZE, OUTPUT_SEQUENCE_SIZE),
    dtype=tf.float32
)
encoder_outputs = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
if RNN_CELL_CLASS._keras_api_names[0] == 'keras.layers.LSTMCell':
    encoder_state = [tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32),
                     tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)]
else:
    # encoder_state = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
    encoder_state = [tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)]

embeddings_dec = keras.layers.Embedding(VOCAB_SIZE, EMBED_SIZE)
decoder_cell = RNN_CELL_CLASS(RNN_UNITS)
output_layer = keras.layers.Dense(VOCAB_SIZE)

sampler = tfa.seq2seq.sampler.TrainingSampler()

decoder = tfa.seq2seq.basic_decoder.BasicDecoder(
    decoder_cell,
    sampler,
    output_layer=output_layer
)

decoder_initial_state = encoder_state

decoder_emb_inp = embeddings_dec(targets)

(
    final_outputs,
    final_state,
    final_sequence_lengths
) = decoder(
    decoder_emb_inp,
    initial_state=decoder_initial_state,
    sequence_length=tf.ones([BATCH_SIZE],
                            dtype=tf.int32) * OUTPUT_SEQUENCE_SIZE
)

```


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
@@ -2046,6 +2046,9 @@ def call(self, inputs, state, **kwargs):
         cell_inputs = self._cell_input_fn(inputs, state.attention)
         cell_state = state.cell_state
         cell_output, next_cell_state = self._cell(cell_inputs, cell_state, **kwargs)
+        next_cell_state = tf.nest.pack_sequence_as(
+            cell_state, tf.nest.flatten(next_cell_state)
+        )
 
         cell_batch_size = (
             tf.compat.dimension_value(cell_output.shape[0]) or tf.shape(cell_output)[0]


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
@@ -920,3 +920,15 @@ def test_attention_state_with_variable_length_input():
     layer = tf.keras.layers.RNN(cell)
 
     _ = layer(data, mask=mask)
+
+
+def test_attention_wrapper_with_gru_cell():
+    mechanism = wrapper.LuongAttention(units=3)
+    cell = tf.keras.layers.GRUCell(3)
+    cell = wrapper.AttentionWrapper(cell, mechanism)
+    memory = tf.ones([2, 5, 3])
+    inputs = tf.ones([2, 3])
+    mechanism.setup_memory(memory)
+    initial_state = cell.get_initial_state(inputs=inputs)
+    _, state = cell(inputs, initial_state)
+    tf.nest.assert_same_structure(initial_state, state)