Repository: tensorflow/addons

Issue:
BeamSearchDecoder with non LSTM cells raises ValueError exception
**System information**
- OS Platform and Distribution: Tensorflow Docker Container based on image tensorflow/tensorflow:2.2.0-jupyter
- TensorFlow version and how it was installed (source or binary): 2.2.0
- TensorFlow-Addons version and how it was installed (source or binary): tfa-nightly  0.11.0.dev20200520032238
- Python version: 3.6.9
- Is GPU used? (yes/no): no

**Describe the bug**

Working with @w4nderlust on this PR https://github.com/uber/ludwig/pull/699  we encountered an issue with seq2seq beam search.  If we use `SimpleRNNCell` or `GRUCell` with beam search and no Atttention we see this error:
```
ValueError: The two structures don't have the same nested structure.
First structure: type=list str=[<tf.Tensor: shape=(384, 256), dtype=float32, numpy=
array([[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.]], dtype=float32)>]
Second structure: type=int str=256
More specifically: Substructure "type=list str=[<tf.Tensor: shape=(384, 256), dtype=float32, numpy=
array([[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.]], dtype=float32)>]" is a sequence, while substructure "type=int str=256" is not

```

No error occurs when we use `LSTMCell` with beam search and no Attention.

Refer this [posting](https://github.com/uber/ludwig/pull/699#issuecomment-636403380) for more context.

**Code to reproduce the issue**
Minimal Working example (**NOTE**: updated minimal working example to simplify code)
``` 
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

#DECODER_CELL_TYPE = keras.layers.SimpleRNNCell  # Raises ValueError Exception
#DECODER_CELL_TYPE = keras.layers.GRUCell        # Raises ValueError Exception
DECODER_CELL_TYPE = keras.layers.LSTMCell        # Does not raise ValueError Exception

VOCAB_SIZE = 100
EMBED_SIZE = 10
RNN_UNITS = 256
INPUT_SEQUENCE_SIZE = 10
OUTPUT_SEQUENCE_SIZE = 15
BATCH_SIZE = 128
BEAM_WIDTH = 3
SIMULATE_LSTM_ENCODER = False

#============== simulate output from encoder ========================
encoder_outputs = tf.zeros([BATCH_SIZE, RNN_UNITS], dtype=tf.float32)
if SIMULATE_LSTM_ENCODER:
    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)]



# ================ Setup Decoder =====================
embeddings_dec = keras.layers.Embedding(VOCAB_SIZE, EMBED_SIZE)
decoder_cell = DECODER_CELL_TYPE(RNN_UNITS)
output_layer = keras.layers.Dense(VOCAB_SIZE)

GO_SYMBOL = VOCAB_SIZE - 1
END_SYMBOL = 0
batch_size = BATCH_SIZE
encoder_sequence_length = tf.convert_to_tensor(
    np.array([INPUT_SEQUENCE_SIZE] * BATCH_SIZE),
    tf.int32
)

decoder_input = tf.expand_dims(
    [GO_SYMBOL] * batch_size, 1)
start_tokens = tf.fill([batch_size], GO_SYMBOL)
end_token = END_SYMBOL

decoder_emb_inp = embeddings_dec(decoder_input)

if DECODER_CELL_TYPE._keras_api_names[0] != 'keras.layers.LSTMCell':
    # adjust for non LSTMCell decoder
    encoder_state = [encoder_state[0]]


#================= setup for beam search ==================
decoder_initial_state = decoder_cell.get_initial_state(
    batch_size=BATCH_SIZE,
    dtype=tf.float32
)

if not isinstance(decoder_initial_state, list):
    decoder_initial_state = [decoder_initial_state]

tiled_decoder_initial_state = tfa.seq2seq.tile_batch(
    decoder_initial_state,
    multiplier=BEAM_WIDTH
)

decoder = tfa.seq2seq.beam_search_decoder.BeamSearchDecoder(
    cell=decoder_cell,
    beam_width=BEAM_WIDTH,
    output_layer=output_layer
)

#================= perform beam search  =====================
decoder_embedding_matrix = embeddings_dec.variables[0]

(
    first_finished,
    first_inputs,
    first_state
) = decoder.initialize(
    decoder_embedding_matrix,
    start_tokens=start_tokens,
    end_token=end_token,
    initial_state=tiled_decoder_initial_state
)

inputs = first_inputs
state = first_state

for j in range(OUTPUT_SEQUENCE_SIZE):
    outputs, next_state, next_inputs, finished = decoder.step(
        j, inputs, state, training=False)
    inputs = next_inputs
    state = next_state

```




Discussion:
@dynamicwebpaige this is another blocker, similar to the similar problem that was solved before.

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py
--- a/tensorflow_addons/seq2seq/beam_search_decoder.py
+++ b/tensorflow_addons/seq2seq/beam_search_decoder.py
@@ -569,6 +569,9 @@ def step(self, time, inputs, state, training=None, name=None):
             cell_outputs = tf.nest.map_structure(
                 lambda out: self._split_batch_beams(out, out.shape[1:]), cell_outputs
             )
+            next_cell_state = tf.nest.pack_sequence_as(
+                cell_state, tf.nest.flatten(next_cell_state)
+            )
             next_cell_state = tf.nest.map_structure(
                 self._maybe_split_batch_beams, next_cell_state, self._cell.state_size
             )


Test patch:
diff --git a/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
--- a/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
+++ b/tensorflow_addons/seq2seq/tests/beam_search_decoder_test.py
@@ -524,7 +524,12 @@ def get_probs():
 @pytest.mark.parametrize("with_alignment_history", [True, False])
 @pytest.mark.parametrize("has_attention", [True, False])
 @pytest.mark.parametrize("time_major", [True, False])
-def test_dynamic_decode_rnn(time_major, has_attention, with_alignment_history):
+@pytest.mark.parametrize(
+    "cell_class", [tf.keras.layers.LSTMCell, tf.keras.layers.GRUCell]
+)
+def test_dynamic_decode_rnn(
+    cell_class, time_major, has_attention, with_alignment_history
+):
     encoder_sequence_length = np.array([3, 2, 3, 1, 1])
     decoder_sequence_length = np.array([2, 0, 1, 2, 3])
     batch_size = 5
@@ -542,7 +547,7 @@ def test_dynamic_decode_rnn(time_major, has_attention, with_alignment_history):
 
     batch_size_tensor = tf.constant(batch_size)
     embedding = np.random.randn(vocab_size, embedding_dim).astype(np.float32)
-    cell = tf.keras.layers.LSTMCell(cell_depth)
+    cell = cell_class(cell_depth)
     initial_state = cell.get_initial_state(batch_size=batch_size, dtype=tf.float32)
     coverage_penalty_weight = 0.0
     if has_attention: