Repository: tensorflow/addons

Issue:
Error when using a mask in RNN + AttentionWrapper cell
**System information**
- OS Platform and Distribution:  manjaro linux testing
- TensorFlow version: pypi tf-nightly 2.2.0.dev20200228
- TensorFlow-Addons version: pypi tfa-nightly 0.9.0.dev20200228
- Python version: 3.7.6
- Is GPU used?: no

**Describe the bug**

A ValueError is raised when using a mask in a RNN layer taking a RNN cell wrapped with an AttentionWrapper.

**Code to reproduce the issue**

```python
import tensorflow_addons as tfa
import tensorflow as tf

bs = 4
cell = tf.keras.layers.LSTMCell(3)
mechanism = tfa.seq2seq.LuongAttention(units=3)
cell = tfa.seq2seq.AttentionWrapper(cell, mechanism)
layer = tf.keras.layers.RNN(cell)


@tf.function
def test(masked=False):
    seq_len = tf.random.uniform(shape=(), minval=2, maxval=10, dtype=tf.int32)
    data = tf.ones(shape=(bs, seq_len, 3))
    mechanism.setup_memory(data)

    if masked:
        mask = tf.sequence_mask(lengths=bs * [seq_len])
    else:
        mask = None

    return layer(data, mask=mask)


test(masked=False); print('Test 1 passed')
test(masked=True); print('Test 2 passed')
```

**Other info / logs**
```python
Test 1 passed
Traceback (most recent call last):
  File "/run/media/zenbook/work/phd/work/__/pb1.py", line 26, in <module>
    test(masked=True); print('Test 2 passed')
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 611, in _call
    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2419, in __call__
    graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2777, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2667, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 981, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 441, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 968, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    /run/media/zenbook/work/phd/work/__/pb1.py:22 test  *
        return layer(data, mask=mask)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/layers/recurrent.py:652 __call__  **
        return super(RNN, self).__call__(inputs, **kwargs)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:926 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/layers/recurrent.py:793 call
        zero_output_for_mask=self.zero_output_for_mask)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/backend.py:4205 rnn
        **while_loop_kwargs)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py:2688 while_loop
        back_prop=back_prop)
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/while_v2.py:229 while_loop
        len_orig_loop_vars], expand_composites=True))
    /home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/while_v2.py:1143 _check_shapes_compat
        "specify a less-specific shape." % (input_t.name, shape, t.shape))

    ValueError: Input tensor 'rnn/AttentionWrapperZeroState/zeros_3:0' enters the loop with shape (), but has shape (4, 1) after one iteration. To allow the shape to vary across iterations, use the `shape_invariants` argument of tf.while_loop to specify a less-specific shape.


Process finished with exit code 1
```



Discussion:
This is because the Keras RNN assumes that each tensor in the state has a batch dimension. However, the `time` field in `AttentionWrapperState` is just a scalar tensor.

I think we could work around not having this time field. @georgesterpu Do you want to look into that?
Thanks, @guillaumekln.
Assuming that the problem comes from the `time` field - `AttentionWrapper` uses it to write the current alignment slice to a `TensorArray` when the `alignment_history` flag is True. I can see several possibilities, such as:
1) pass the current loop index from outside (the `time` parameter of the `_step` function declared inside `keras.backend.rnn` ?
2) manage the loop counter as a member of the `AttentionWrapper` class
3) extend the `TensorArray` class to support an `append` operation, since `AttentionWrapper` never leapfrogs over indices and only writes sequentially anyway. 

Unfortunately, as a final year student, I lack the throughput to carry on with a more detailed investigation - this is why I created the issue in the first place :relaxed: 
Thanks looking into workarounds. Actually, as `alignment_history` has a dynamic size and we only append to it, we can do something like:

```python
ta = ta.write(ta.size(), elem)
```

> Unfortunately, as a final year student, I lack the throughput to carry on with a more detailed investigation - this is why I created the issue in the first place relaxed

No worries! Thank you for finding these issues. I can make a PR as I already have a fix locally.

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
@@ -1389,7 +1389,6 @@ class AttentionWrapperState(
         (
             "cell_state",
             "attention",
-            "time",
             "alignments",
             "alignment_history",
             "attention_state",
@@ -1403,7 +1402,6 @@ class AttentionWrapperState(
       - `cell_state`: The state of the wrapped `RNNCell` at the previous time
         step.
       - `attention`: The attention emitted at the previous time step.
-      - `time`: int32 scalar containing the current time step.
       - `alignments`: A single or tuple of `Tensor`(s) containing the
          alignments emitted at the previous time step for each attention
          mechanism.
@@ -1915,7 +1913,6 @@ def state_size(self):
         """
         return AttentionWrapperState(
             cell_state=self._cell.state_size,
-            time=tf.TensorShape([]),
             attention=self._get_attention_layer_size(),
             alignments=self._item_or_tuple(
                 a.alignments_size for a in self._attention_mechanisms
@@ -1983,7 +1980,6 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
             ]
             return AttentionWrapperState(
                 cell_state=cell_state,
-                time=tf.zeros([], dtype=tf.int32),
                 attention=tf.zeros(
                     [batch_size, self._get_attention_layer_size()], dtype=dtype
                 ),
@@ -2086,7 +2082,9 @@ def call(self, inputs, state, **kwargs):
                 self._attention_layers[i] if self._attention_layers else None,
             )
             alignment_history = (
-                previous_alignment_history[i].write(state.time, alignments)
+                previous_alignment_history[i].write(
+                    previous_alignment_history[i].size(), alignments
+                )
                 if self._alignment_history
                 else ()
             )
@@ -2098,7 +2096,6 @@ def call(self, inputs, state, **kwargs):
 
         attention = tf.concat(all_attentions, 1)
         next_state = AttentionWrapperState(
-            time=state.time + 1,
             cell_state=next_cell_state,
             attention=attention,
             attention_state=self._item_or_tuple(all_attention_states),


Test patch:
diff --git a/tensorflow_addons/seq2seq/attention_wrapper_test.py b/tensorflow_addons/seq2seq/attention_wrapper_test.py
--- a/tensorflow_addons/seq2seq/attention_wrapper_test.py
+++ b/tensorflow_addons/seq2seq/attention_wrapper_test.py
@@ -476,7 +476,7 @@ def _testWithMaybeMultiAttention(
             self.assertIsInstance(final_state, wrapper.AttentionWrapperState)
 
             expected_time = (
-                expected_final_state.time if tf.executing_eagerly() else None
+                max(decoder_sequence_length) if tf.executing_eagerly() else None
             )
             self.assertEqual(
                 (batch_size, expected_time, attention_depth),
@@ -645,7 +645,6 @@ def testBahdanauNotNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype(np.float32), mean=0.041453815
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype(np.float32), mean=0.125
             ),
@@ -686,7 +685,6 @@ def testBahdanauNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=0.042427111
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.125
             ),
@@ -723,7 +721,6 @@ def testLuongNotNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=-0.0318060
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.125
             ),
@@ -760,7 +757,6 @@ def testLuongScaled(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=-0.0318060
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.125
             ),
@@ -796,7 +792,6 @@ def testNotUseAttentionLayer(self):
             attention=ResultSummary(
                 shape=(5, 10), dtype=np.dtype("float32"), mean=0.026356646
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.125
             ),
@@ -835,7 +830,6 @@ def testBahdanauMonotonicNotNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=0.038682378
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.09778417
             ),
@@ -877,7 +871,6 @@ def testBahdanauMonotonicNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=0.068432882
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.0615656
             ),
@@ -919,7 +912,6 @@ def testLuongMonotonicNotNormalized(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=0.059128221
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.05112994
             ),
@@ -961,7 +953,6 @@ def testLuongMonotonicScaled(self):
             attention=ResultSummary(
                 shape=(5, 6), dtype=np.dtype("float32"), mean=0.059128221
             ),
-            time=3,
             alignments=ResultSummary(
                 shape=(5, 8), dtype=np.dtype("float32"), mean=0.05112994
             ),
@@ -1005,12 +996,16 @@ def test_attention_state_with_variable_length_input(self):
         cell = wrapper.AttentionWrapper(cell, mechanism)
 
         var_len = tf.random.uniform(shape=(), minval=2, maxval=10, dtype=tf.int32)
+        lengths = tf.random.uniform(
+            shape=(var_len,), minval=1, maxval=var_len + 1, dtype=tf.int32
+        )
         data = tf.ones(shape=(var_len, var_len, 3))
+        mask = tf.sequence_mask(lengths, maxlen=var_len)
 
         mechanism.setup_memory(data)
         layer = tf.keras.layers.RNN(cell)
 
-        _ = layer(data)
+        _ = layer(data, mask=mask)
 
 
 if __name__ == "__main__":