Accept passing maximum_iterations as a Tensor (#1751)

Also add `@typechecked` to dynamic_decode for consistency.

Author: guillaumekln

tensorflow_addons/seq2seq/decoder.py

@@ -17,7 +17,7 @@
 import abc
 
 import tensorflow as tf
-from tensorflow_addons.utils.types import Number
+from tensorflow_addons.utils.types import TensorLike
 from typeguard import typechecked
 from typing import Any, Optional, Tuple, Union
 
@@ -144,7 +144,7 @@ def __init__(
         self,
         output_time_major: bool = False,
         impute_finished: bool = False,
-        maximum_iterations: Optional[Number] = None,
+        maximum_iterations: Optional[TensorLike] = None,
         parallel_iterations: int = 32,
         swap_memory: bool = False,
         **kwargs
@@ -256,11 +256,12 @@ def tracks_own_finished(self):
     # TODO(scottzhu): Add build/get_config/from_config and other layer methods.
 
 
+@typechecked
 def dynamic_decode(
     decoder: Union[Decoder, BaseDecoder],
     output_time_major: bool = False,
     impute_finished: bool = False,
-    maximum_iterations: Optional[Number] = None,
+    maximum_iterations: Optional[TensorLike] = None,
     parallel_iterations: int = 32,
     swap_memory: bool = False,
     training: Optional[bool] = None,
@@ -299,14 +300,8 @@ def dynamic_decode(
       `(final_outputs, final_state, final_sequence_lengths)`.
 
     Raises:
-      TypeError: if `decoder` is not an instance of `Decoder`.
       ValueError: if `maximum_iterations` is provided but is not a scalar.
     """
-    if not isinstance(decoder, (Decoder, BaseDecoder)):
-        raise TypeError(
-            "Expected decoder to be type Decoder, but saw: %s" % type(decoder)
-        )
-
     with tf.compat.v1.variable_scope(scope, "decoder") as varscope:
         # Determine context types.
         ctxt = tf.compat.v1.get_default_graph()._get_control_flow_context()

tensorflow_addons/seq2seq/tests/decoder_test.py

@@ -23,7 +23,9 @@
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
-@pytest.mark.parametrize("maximum_iterations", [None, 0, 1])
+@pytest.mark.parametrize(
+    "maximum_iterations", [None, 0, 1, tf.constant(1, dtype=tf.int32)]
+)
 @pytest.mark.parametrize("time_major", [True, False])
 def test_dynamic_decode_rnn(time_major, maximum_iterations):