Repository: tensorflow/addons

Issue:
Incorrect type of maximum_iterations argument in tfa.seq2seq.decoder causes runtime error.
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Debian Stable
- TensorFlow version and how it was installed (source or binary): 2.1 binary
- TensorFlow-Addons version and how it was installed (source or binary): 0.8.1 binary
- Python version: 3.7
- Is GPU used? (yes/no): no

**Describe the bug**

Calling `tfa.seq2seq.decoder.BaseDecoder` with a `tf.Tensor` for `maximum_iterations`, which worked before, now results in

```
TypeError: type of argument "maximum_iterations" must be one of (float, int, float16, float32, int8, int16, int32, int64, uint8, uint16, uint32, uint64, NoneType); got tensorflow.python.framework.ops.Tensor instead
```

**Other info / logs**

The problem is caused by a faulty type annotation
```python
  maximum_iterations: Optional[Number] = None,
```

Note that this annotation is also used in `tfa.seq2seq.decoder.dynamic_decode`.


Discussion:
Note that while `tfa.seq2seq.decoder.BaseDecoder.__init__` is marked with `@typechecked`, `tfa.seq2seq.decoder.dynamic_decode` is not.

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/decoder.py b/tensorflow_addons/seq2seq/decoder.py
--- a/tensorflow_addons/seq2seq/decoder.py
+++ b/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()


Test patch:
diff --git a/tensorflow_addons/seq2seq/tests/decoder_test.py b/tensorflow_addons/seq2seq/tests/decoder_test.py
--- a/tensorflow_addons/seq2seq/tests/decoder_test.py
+++ b/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):