Repository: tensorflow/addons

Issue:
CRF functions fail on sequences of length 1 when the static dimension is undefined
**System information**
- OS Platform and Distribution: Ubuntu 16.04
- TensorFlow version and how it was installed: 2.0.0 (binary)
- TensorFlow-Addons version and how it was installed: 0.6.0 (binary)
- Python version: 3.5.2
- Is GPU used?: No

**Describe the bug**

Some CRF functions have a special case for sequence length = 1. However, this code path is not enabled when the static dimension is undefined and a `InvalidArgumentError` exception is raised. It can happen when CRF functions are used in `tf.function`.

Conditions like:

```python
if inputs.shape[1] == 1:
```

should probably be:

```python
if tf.shape(inputs)[1] == 1:
```

**Code to reproduce the issue**

```python
import tensorflow as tf
import tensorflow_addons as tfa

num_tags = 10
input_signature = (
    tf.TensorSpec([None, None, num_tags]),
    tf.TensorSpec([num_tags, num_tags]),
    tf.TensorSpec([None], dtype=tf.int32))
crf_decode = tf.function(input_signature=input_signature)(tfa.text.crf_decode)

batch_size = 4
potentials = tf.random.uniform([batch_size, 1, num_tags])
transition_params = tf.random.uniform([num_tags, num_tags])
sequence_length = tf.ones([batch_size], dtype=tf.int32)

crf_decode(potentials, transition_params, sequence_length)
```

**Other info / logs**

```text
2019-11-12 14:17:08.541583: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: Tried to stack elements of an empty list with non-fully-defined element_shape: [?,10]
         [[{{node rnn/TensorArrayV2Stack/TensorListStack}}]]
         [[concat/_30]]
2019-11-12 14:17:08.541618: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: Tried to stack elements of an empty list with non-fully-defined element_shape: [?,10]
         [[{{node rnn/TensorArrayV2Stack/TensorListStack}}]]
Traceback (most recent call last):
  File "test/crf.py", line 15, in <module>
    crf_decode(potentials, transition_params, sequence_length)
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/def_function.py", line 526, in _call
    return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)  # pylint: disable=protected-access
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/function.py", line 1141, in _filtered_call
    self.captured_inputs)
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/function.py", line 1224, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/function.py", line 511, in call
    ctx=ctx)
  File "lib/python3.5/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Tried to stack elements of an empty list with non-fully-defined element_shape: [?,10]
         [[node rnn/TensorArrayV2Stack/TensorListStack (defined at lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:1751) ]]
         [[concat/_30]]
  (1) Invalid argument:  Tried to stack elements of an empty list with non-fully-defined element_shape: [?,10]
         [[node rnn/TensorArrayV2Stack/TensorListStack (defined at lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:1751) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_crf_decode_264]

Function call stack:
crf_decode -> crf_decode
```


Discussion:

Resolving patch:
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -44,13 +44,16 @@ def crf_sequence_score(inputs, tag_indices, sequence_lengths,
     # If max_seq_len is 1, we skip the score calculation and simply gather the
     # unary potentials of the single tag.
     def _single_seq_fn():
-        batch_size = tf.shape(inputs, out_type=tag_indices.dtype)[0]
+        batch_size = tf.shape(inputs, out_type=tf.int32)[0]
+        batch_inds = tf.reshape(tf.range(batch_size), [-1, 1])
+        indices = tf.concat([batch_inds, tf.zeros_like(batch_inds)], axis=1)
+
+        tag_inds = tf.gather_nd(tag_indices, indices)
+        tag_inds = tf.reshape(tag_inds, [-1, 1])
+        indices = tf.concat([indices, tag_inds], axis=1)
+
+        sequence_scores = tf.gather_nd(inputs, indices)
 
-        example_inds = tf.reshape(
-            tf.range(batch_size, dtype=tag_indices.dtype), [-1, 1])
-        sequence_scores = tf.gather_nd(
-            tf.squeeze(inputs, [1]),
-            tf.concat([example_inds, tag_indices], axis=1))
         sequence_scores = tf.where(
             tf.less_equal(sequence_lengths, 0), tf.zeros_like(sequence_scores),
             sequence_scores)
@@ -64,10 +67,8 @@ def _multi_seq_fn():
         sequence_scores = unary_scores + binary_scores
         return sequence_scores
 
-    if inputs.shape[1] == 1:
-        return _single_seq_fn()
-    else:
-        return _multi_seq_fn()
+    return tf.cond(
+        tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
 
 
 def crf_multitag_sequence_score(inputs, tag_bitmap, sequence_lengths,
@@ -106,16 +107,15 @@ def _single_seq_fn():
             filtered_inputs, axis=[1, 2], keepdims=False)
 
     def _multi_seq_fn():
-        # Compute the logsumexp of all scores of sequences matching the given tags.
+        # Compute the logsumexp of all scores of sequences
+        # matching the given tags.
         return crf_log_norm(
             inputs=filtered_inputs,
             sequence_lengths=sequence_lengths,
             transition_params=transition_params)
 
-    if inputs.shape[1] == 1:
-        return _single_seq_fn()
-    else:
-        return _multi_seq_fn()
+    return tf.cond(
+        tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
 
 
 def crf_log_norm(inputs, sequence_lengths, transition_params):
@@ -135,8 +135,8 @@ def crf_log_norm(inputs, sequence_lengths, transition_params):
     first_input = tf.slice(inputs, [0, 0, 0], [-1, 1, -1])
     first_input = tf.squeeze(first_input, [1])
 
-    # If max_seq_len is 1, we skip the algorithm and simply reduce_logsumexp over
-    # the "initial state" (the unary potentials).
+    # If max_seq_len is 1, we skip the algorithm and simply reduce_logsumexp
+    # over the "initial state" (the unary potentials).
     def _single_seq_fn():
         log_norm = tf.reduce_logsumexp(first_input, [1])
         # Mask `log_norm` of the sequences with length <= zero.
@@ -160,10 +160,8 @@ def _multi_seq_fn():
             log_norm)
         return log_norm
 
-    if inputs.shape[1] == 1:
-        return _single_seq_fn()
-    else:
-        return _multi_seq_fn()
+    return tf.cond(
+        tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
 
 
 def crf_log_likelihood(inputs,
@@ -444,9 +442,7 @@ def _scan_fn(state, inputs):
 
 
 def crf_decode(potentials, transition_params, sequence_length):
-    """Decode the highest scoring sequence of tags in TensorFlow.
-
-    This is a function for tensor.
+    """Decode the highest scoring sequence of tags.
 
     Args:
       potentials: A [batch_size, max_seq_len, num_tags] tensor of
@@ -462,23 +458,21 @@ def crf_decode(potentials, transition_params, sequence_length):
     """
     sequence_length = tf.cast(sequence_length, dtype=tf.int32)
 
-    # If max_seq_len is 1, we skip the algorithm and simply return the argmax tag
-    # and the max activation.
+    # If max_seq_len is 1, we skip the algorithm and simply return the
+    # argmax tag and the max activation.
     def _single_seq_fn():
-        squeezed_potentials = tf.squeeze(potentials, [1])
-        decode_tags = tf.expand_dims(tf.argmax(squeezed_potentials, axis=1), 1)
-        best_score = tf.reduce_max(squeezed_potentials, axis=1)
-        return tf.cast(decode_tags, dtype=tf.int32), best_score
+        decode_tags = tf.cast(tf.argmax(potentials, axis=2), dtype=tf.int32)
+        best_score = tf.reshape(tf.reduce_max(potentials, axis=2), shape=[-1])
+        return decode_tags, best_score
 
     def _multi_seq_fn():
-        """Decoding of highest scoring sequence."""
         # Computes forward decoding. Get last score and backpointers.
         initial_state = tf.slice(potentials, [0, 0, 0], [-1, 1, -1])
         initial_state = tf.squeeze(initial_state, axis=[1])
         inputs = tf.slice(potentials, [0, 1, 0], [-1, -1, -1])
 
         sequence_length_less_one = tf.maximum(
-            tf.constant(0, dtype=sequence_length.dtype), sequence_length - 1)
+            tf.constant(0, dtype=tf.int32), sequence_length - 1)
 
         backpointers, last_score = crf_decode_forward(
             inputs, initial_state, transition_params, sequence_length_less_one)
@@ -498,7 +492,14 @@ def _multi_seq_fn():
         best_score = tf.reduce_max(last_score, axis=1)
         return decode_tags, best_score
 
-    if potentials.shape[1] == 1:
-        return _single_seq_fn()
+    if potentials.shape[1] is not None:
+        # shape is statically know, so we just execute
+        # the appropriate code path
+        if potentials.shape[1] == 1:
+            return _single_seq_fn()
+        else:
+            return _multi_seq_fn()
     else:
-        return _multi_seq_fn()
+        return tf.cond(
+            tf.equal(tf.shape(potentials)[1], 1), _single_seq_fn,
+            _multi_seq_fn)


Test patch:
diff --git a/tensorflow_addons/text/crf_test.py b/tensorflow_addons/text/crf_test.py
--- a/tensorflow_addons/text/crf_test.py
+++ b/tensorflow_addons/text/crf_test.py
@@ -356,13 +356,24 @@ def testCrfDecodeZeroSeqLength(self):
     def testDifferentDtype(self):
         inputs = np.ones([16, 20, 5], dtype=np.float32)
         tags = tf.convert_to_tensor(np.ones([16, 20], dtype=np.int64))
-        seq_lens = np.ones([
-            16,
-        ], dtype=np.int64) * 20
+        seq_lens = np.ones([16], dtype=np.int64) * 20
 
         loss, _ = text.crf_log_likelihood(
             inputs=inputs, tag_indices=tags, sequence_lengths=seq_lens)
 
+    def testTfFunction(self):
+        batch_size = 4
+        num_tags = 10
+        input_signature = (tf.TensorSpec([None, None, num_tags]),
+                           tf.TensorSpec([num_tags, num_tags]),
+                           tf.TensorSpec([None], dtype=tf.int32))
+        crf_decode = tf.function(input_signature=input_signature)(
+            text.crf_decode)
+        potentials = tf.random.uniform([batch_size, 1, num_tags])
+        transition_params = tf.random.uniform([num_tags, num_tags])
+        sequence_length = tf.ones([batch_size], dtype=tf.int32)
+        crf_decode(potentials, transition_params, sequence_length)
+
 
 if __name__ == "__main__":
     tf.test.main()