fix state size to be a TensorShape instead of a Tensor (#1121)

* fix state size to be a TensorShape instead of a Tensor

* fix Eager mode tests where alignments_size is an int

* fix 2

* add a test case

* update module name

* create an empty TensorShape when self._alignments_size is a Tensor

Co-authored-by: Gabriel de Marmiesse <[email protected]>

Author: georgesterpu

tensorflow_addons/seq2seq/attention_wrapper.py

@@ -411,11 +411,14 @@ def deserialize_inner_layer_from_config(cls, config, custom_objects):
 
     @property
     def alignments_size(self):
-        return self._alignments_size
+        if isinstance(self._alignments_size, int):
+            return self._alignments_size
+        else:
+            return tf.TensorShape([None])
 
     @property
     def state_size(self):
-        return self._alignments_size
+        return self.alignments_size
 
     def initial_alignments(self, batch_size, dtype):
         """Creates the initial alignment values for the `AttentionWrapper`

tensorflow_addons/seq2seq/attention_wrapper_test.py

@@ -999,6 +999,19 @@ def test_attention_state_with_keras_rnn(self):
         initial_state = cell.get_initial_state(batch_size=2, dtype=tf.float32)
         _ = layer(inputs=tf.ones((2, 4, 8)), initial_state=initial_state)
 
+    def test_attention_state_with_variable_length_input(self):
+        cell = tf.keras.layers.LSTMCell(3)
+        mechanism = wrapper.LuongAttention(units=3)
+        cell = wrapper.AttentionWrapper(cell, mechanism)
+
+        var_len = tf.random.uniform(shape=(), minval=2, maxval=10, dtype=tf.int32)
+        data = tf.ones(shape=(var_len, var_len, 3))
+
+        mechanism.setup_memory(data)
+        layer = tf.keras.layers.RNN(cell)
+
+        _ = layer(data)
+
 
 if __name__ == "__main__":
     tf.test.main()