Repository: tensorflow/addons

Issue:
Make tfa.seq2seq.AttentionWrapper 2.0 compatible
The `AttentionWrapper` class is currently not compatible with TF2.0:

* it extends `tf.compat.v1.nn.rnn_cell.RNNCell` instead of `tf.keras.layers.AbstractRNNCell`
* it uses `tf.compat.v1.layers.Dense` instead of `tf.keras.layers.Dense`
* attention layers are saved in a `tuple` instead of a `list` which prevent them to be tracked

I can contribute to this change, unless this is a WIP by @qlzh727.


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
@@ -31,7 +31,6 @@
 from tensorflow.python.keras import initializers
 from tensorflow.python.keras import layers
 from tensorflow.python.keras.engine import base_layer_utils
-from tensorflow.python.layers import core as layers_core
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import clip_ops
@@ -1369,7 +1368,8 @@ def clone(self, **kwargs):
         Example:
 
         ```python
-        initial_state = attention_wrapper.zero_state(dtype=..., batch_size=...)
+        initial_state = attention_wrapper.get_initial_state(
+            batch_size=..., dtype=...)
         initial_state = initial_state.clone(cell_state=encoder_state)
         ```
 
@@ -1558,7 +1558,7 @@ def _compute_attention(attention_mechanism, cell_output, attention_state,
     return attention, alignments, next_attention_state
 
 
-class AttentionWrapper(rnn_cell_impl.RNNCell):
+class AttentionWrapper(layers.AbstractRNNCell):
     """Wraps another `RNNCell` with attention."""
 
     def __init__(self,
@@ -1579,9 +1579,9 @@ def __init__(self,
 
         - The encoder output has been tiled to `beam_width` via
           `tf.contrib.seq2seq.tile_batch` (NOT `tf.tile`).
-        - The `batch_size` argument passed to the `zero_state` method of this
-          wrapper is equal to `true_batch_size * beam_width`.
-        - The initial state created with `zero_state` above contains a
+        - The `batch_size` argument passed to the `get_initial_state` method of
+          this wrapper is equal to `true_batch_size * beam_width`.
+        - The initial state created with `get_initial_state` above contains a
           `cell_state` value containing properly tiled final state from the
           encoder.
 
@@ -1599,8 +1599,8 @@ def __init__(self,
             memory=tiled_inputs,
             memory_sequence_length=tiled_sequence_length)
         attention_cell = AttentionWrapper(cell, attention_mechanism, ...)
-        decoder_initial_state = attention_cell.zero_state(
-            dtype, batch_size=true_batch_size * beam_width)
+        decoder_initial_state = attention_cell.get_initial_state(
+            batch_size=true_batch_size * beam_width, dtype=dtype)
         decoder_initial_state = decoder_initial_state.clone(
             cell_state=tiled_encoder_final_state)
         ```
@@ -1634,10 +1634,10 @@ def __init__(self,
             attention mechanism is propagated up to the next cell in an RNN
             stack or to the top RNN output.
           initial_cell_state: The initial state value to use for the cell when
-            the user calls `zero_state()`.  Note that if this value is provided
-            now, and the user uses a `batch_size` argument of `zero_state`
-            which does not match the batch size of `initial_cell_state`, proper
-            behavior is not guaranteed.
+            the user calls `get_initial_state()`.  Note that if this value is
+            provided now, and the user uses a `batch_size` argument of
+            `get_initial_state` which does not match the batch size of
+            `initial_cell_state`, proper behavior is not guaranteed.
           name: Name to use when creating ops.
           attention_layer: A list of `tf.layers.Layer` instances or a
             single `tf.layers.Layer` instance taking the context and cell
@@ -1666,7 +1666,7 @@ def __init__(self,
         rnn_cell_impl.assert_like_rnncell("cell", cell)
         if isinstance(attention_mechanism, (list, tuple)):
             self._is_multi = True
-            attention_mechanisms = attention_mechanism
+            attention_mechanisms = list(attention_mechanism)
             for attention_mechanism in attention_mechanisms:
                 if not isinstance(attention_mechanism, AttentionMechanism):
                     raise TypeError(
@@ -1680,7 +1680,7 @@ def __init__(self,
                     "attention_mechanism must be an AttentionMechanism or "
                     "list of multiple AttentionMechanism instances, saw type: "
                     "%s" % type(attention_mechanism).__name__)
-            attention_mechanisms = (attention_mechanism,)
+            attention_mechanisms = [attention_mechanism]
 
         if cell_input_fn is None:
             cell_input_fn = (lambda inputs, attention: array_ops.concat(
@@ -1704,8 +1704,8 @@ def __init__(self,
                     "If provided, attention_layer_size must contain exactly "
                     "one integer per attention_mechanism, saw: %d vs %d" %
                     (len(attention_layer_sizes), len(attention_mechanisms)))
-            self._attention_layers = tuple(
-                layers_core.Dense(
+            self._attention_layers = list(
+                layers.Dense(
                     attention_layer_size,
                     name="attention_layer",
                     use_bias=False,
@@ -1713,7 +1713,7 @@ def __init__(self,
                 attention_layer_size in enumerate(attention_layer_sizes))
             self._attention_layer_size = sum(attention_layer_sizes)
         elif attention_layer is not None:
-            self._attention_layers = tuple(
+            self._attention_layers = list(
                 attention_layer if isinstance(attention_layer, (
                     list, tuple)) else (attention_layer,))
             if len(self._attention_layers) != len(attention_mechanisms):
@@ -1753,8 +1753,8 @@ def __init__(self,
                     final_state_tensor.shape[0])
                                     or array_ops.shape(final_state_tensor)[0])
                 error_message = (
-                    "When constructing AttentionWrapper %s: " % self._base_name
-                    + "Non-matching batch sizes between the memory "
+                    "When constructing AttentionWrapper %s: " % self.name +
+                    "Non-matching batch sizes between the memory "
                     "(encoder output) and initial_cell_state.  Are you using "
                     "the BeamSearchDecoder?  You may need to tile your "
                     "initial state via the tf.contrib.seq2seq.tile_batch "
@@ -1823,14 +1823,15 @@ def state_size(self):
                 a.alignments_size if self._alignment_history else () for a in
                 self._attention_mechanisms))  # sometimes a TensorArray
 
-    def zero_state(self, batch_size, dtype):
+    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
         """Return an initial (zero) state tuple for this `AttentionWrapper`.
 
         **NOTE** Please see the initializer documentation for details of how
-        to call `zero_state` if using an `AttentionWrapper` with a
+        to call `get_initial_state` if using an `AttentionWrapper` with a
         `BeamSearchDecoder`.
 
         Args:
+          inputs: The inputs that will be fed to this cell.
           batch_size: `0D` integer tensor: the batch size.
           dtype: The internal state data type.
 
@@ -1843,6 +1844,9 @@ def zero_state(self, batch_size, dtype):
             `batch_size` does not match the output size of the encoder passed
             to the wrapper object at initialization time.
         """
+        if inputs is not None:
+            batch_size = array_ops.shape(inputs)[0]
+            dtype = inputs.dtype
         with ops.name_scope(
                 type(self).__name__ + "ZeroState", values=[batch_size]):  # pylint: disable=bad-continuation
             if self._initial_cell_state is not None:
@@ -1851,14 +1855,13 @@ def zero_state(self, batch_size, dtype):
                 cell_state = self._cell.get_initial_state(
                     batch_size=batch_size, dtype=dtype)
             error_message = (
-                "When calling zero_state of AttentionWrapper %s: " %
-                self._base_name +
-                "Non-matching batch sizes between the memory "
+                "When calling get_initial_state of AttentionWrapper %s: " %
+                self.name + "Non-matching batch sizes between the memory "
                 "(encoder output) and the requested batch size. Are you using "
                 "the BeamSearchDecoder?  If so, make sure your encoder output "
                 "has been tiled to beam_width via "
                 "tf.contrib.seq2seq.tile_batch, and the batch_size= argument "
-                "passed to zero_state is batch_size * beam_width.")
+                "passed to get_initial_state is batch_size * beam_width.")
             with ops.control_dependencies(
                     self._batch_size_checks(batch_size, error_message)):  # pylint: disable=bad-continuation
                 cell_state = nest.map_structure(
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
@@ -450,7 +450,7 @@ def _split_batch_beams(self, t, s=None):
                 "and batch size.  The reshaped tensor has shape: %s.  "
                 "We expected it to have shape "
                 "(batch_size, beam_width, depth) == %s.  Perhaps you "
-                "forgot to create a zero_state with "
+                "forgot to call get_initial_state with "
                 "batch_size=encoder_batch_size * beam_width?" %
                 (reshaped_t.shape, expected_reshaped_shape))
         reshaped_t.set_shape(expected_reshaped_shape)
@@ -622,9 +622,9 @@ class BeamSearchDecoder(BeamSearchDecoderMixin, decoder.BaseDecoder):
 
     - The encoder output has been tiled to `beam_width` via
       `tf.contrib.seq2seq.tile_batch` (NOT `tf.tile`).
-    - The `batch_size` argument passed to the `zero_state` method of this
-      wrapper is equal to `true_batch_size * beam_width`.
-    - The initial state created with `zero_state` above contains a
+    - The `batch_size` argument passed to the `get_initial_state` method of
+      this wrapper is equal to `true_batch_size * beam_width`.
+    - The initial state created with `get_initial_state` above contains a
       `cell_state` value containing properly tiled final state from the
       encoder.
 
@@ -642,8 +642,8 @@ class BeamSearchDecoder(BeamSearchDecoderMixin, decoder.BaseDecoder):
         memory=tiled_inputs,
         memory_sequence_length=tiled_sequence_length)
     attention_cell = AttentionWrapper(cell, attention_mechanism, ...)
-    decoder_initial_state = attention_cell.zero_state(
-        dtype, batch_size=true_batch_size * beam_width)
+    decoder_initial_state = attention_cell.get_initial_state(
+        batch_size=true_batch_size * beam_width, dtype=dtype)
     decoder_initial_state = decoder_initial_state.clone(
         cell_state=tiled_encoder_final_state)
     ```


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
@@ -319,9 +319,8 @@ def _testWithMaybeMultiAttention(self,
                 attention_layer=attention_layer)
             if cell._attention_layers is not None:
                 for layer in cell._attention_layers:
-                    if getattr(layer, "kernel_initializer") is None:
-                        layer.kernel_initializer = initializers.glorot_uniform(
-                            seed=1337)
+                    layer.kernel_initializer = initializers.glorot_uniform(
+                        seed=1337)
 
             sampler = sampler_py.TrainingSampler()
             my_decoder = basic_decoder.BasicDecoder(cell=cell, sampler=sampler)
@@ -372,8 +371,9 @@ def _testWithMaybeMultiAttention(self,
                         (expected_time, batch_size, encoder_max_time),
                         tuple(state_alignment_history.get_shape().as_list()))
                 nest.assert_same_structure(
-                    cell.state_size, cell.zero_state(batch_size,
-                                                     dtypes.float32))
+                    cell.state_size,
+                    cell.get_initial_state(
+                        batch_size=batch_size, dtype=dtypes.float32))
                 # Remove the history from final_state for purposes of the
                 # remainder of the tests.
                 final_state = final_state._replace(alignment_history=())  # pylint: disable=protected-access
@@ -433,7 +433,8 @@ def testBahdanauNormalizedDType(self, dtype):
 
         final_outputs, final_state, _ = my_decoder(
             decoder_inputs,
-            initial_state=cell.zero_state(dtype=dtype, batch_size=self.batch),
+            initial_state=cell.get_initial_state(
+                batch_size=self.batch, dtype=dtype),
             sequence_length=self.decoder_sequence_length)
         self.assertIsInstance(final_outputs, basic_decoder.BasicDecoderOutput)
         self.assertEqual(final_outputs.rnn_output.dtype, dtype)
@@ -460,7 +461,8 @@ def testLuongScaledDType(self, dtype):
 
         final_outputs, final_state, _ = my_decoder(
             decoder_inputs,
-            initial_state=cell.zero_state(dtype=dtype, batch_size=self.batch),
+            initial_state=cell.get_initial_state(
+                batch_size=self.batch, dtype=dtype),
             sequence_length=self.decoder_sequence_length)
         self.assertIsInstance(final_outputs, basic_decoder.BasicDecoderOutput)
         self.assertEqual(final_outputs.rnn_output.dtype, dtype)
diff --git a/tensorflow_addons/seq2seq/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/beam_search_decoder_test.py
--- a/tensorflow_addons/seq2seq/beam_search_decoder_test.py
+++ b/tensorflow_addons/seq2seq/beam_search_decoder_test.py
@@ -30,7 +30,6 @@
 from tensorflow.python.keras import layers
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import nn_ops
-from tensorflow.python.ops import rnn_cell
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
@@ -508,8 +507,9 @@ def _testDynamicDecodeRNN(self,
             batch_size_tensor = constant_op.constant(batch_size)
             embedding = np.random.randn(vocab_size,
                                         embedding_dim).astype(np.float32)
-            cell = rnn_cell.LSTMCell(cell_depth)
-            initial_state = cell.zero_state(batch_size, dtypes.float32)
+            cell = layers.LSTMCell(cell_depth)
+            initial_state = cell.get_initial_state(
+                batch_size=batch_size, dtype=dtypes.float32)
             coverage_penalty_weight = 0.0
             if has_attention:
                 coverage_penalty_weight = 0.2
@@ -532,9 +532,9 @@ def _testDynamicDecodeRNN(self,
                     attention_mechanism=attention_mechanism,
                     attention_layer_size=attention_depth,
                     alignment_history=with_alignment_history)
-            cell_state = cell.zero_state(
-                dtype=dtypes.float32,
-                batch_size=batch_size_tensor * beam_width)
+            cell_state = cell.get_initial_state(
+                batch_size=batch_size_tensor * beam_width,
+                dtype=dtypes.float32)
             if has_attention:
                 cell_state = cell_state.clone(cell_state=initial_state)
             bsd = beam_search_decoder.BeamSearchDecoder(