Repository: tensorflow/addons

Issue:
InvalidArgumentError caused by int32 / int64 in crf_log_likelihood
**System information**
- OS: Ubuntu 18.04
- TensorFlow: 2.0.0, pip install tensorflow-gpu
- TensorFlow-Addons: 0.6.0, pip install tensorflow-addons
- Python version: 3.6.7
- Is GPU used? (yes/no): yes
- CUDA: 10.0
- CUDNN: 7.6.4

**Describe the bug**

tf.shape returns int32 but Mul op expects int64 (in ```crf_binary_score``` function)
tf.maximum returns int64 but Pack op expects int32 (in ```crf_forward``` function)

**Code to reproduce the issue**

```
import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np

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
loss, _ = tfa.text.crf_log_likelihood(
    inputs=inputs,
    tag_indices=tags,
    sequence_lengths=seq_lens
)
```

**Other info / logs**

Error logs:
```
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-2-41e7d4eaa2ab> in <module>
      6     inputs=inputs,
      7     tag_indices=tags,
----> 8     sequence_lengths=seq_lens
      9 )

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_log_likelihood(inputs, tag_indices, sequence_lengths, transition_params)
    194 
    195     sequence_scores = crf_sequence_score(inputs, tag_indices, sequence_lengths,
--> 196                                          transition_params)
    197     log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)
    198 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_sequence_score(inputs, tag_indices, sequence_lengths, transition_params)
     66         return _single_seq_fn()
     67     else:
---> 68         return _multi_seq_fn()
     69 
     70 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in _multi_seq_fn()
     59         unary_scores = crf_unary_score(tag_indices, sequence_lengths, inputs)
     60         binary_scores = crf_binary_score(tag_indices, sequence_lengths,
---> 61                                          transition_params)
     62         sequence_scores = unary_scores + binary_scores
     63         return sequence_scores

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_binary_score(tag_indices, sequence_lengths, transition_params)
    257     # Encode the indices in a flattened representation.
    258     flattened_transition_indices = start_tag_indices * \
--> 259         num_tags + end_tag_indices
    260     flattened_transition_params = tf.reshape(transition_params, [-1])
    261 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/ops/math_ops.py in binary_op_wrapper(x, y)
    897     with ops.name_scope(None, op_name, [x, y]) as name:
    898       if isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor):
--> 899         return func(x, y, name=name)
    900       elif not isinstance(y, sparse_tensor.SparseTensor):
    901         try:

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/ops/math_ops.py in _mul_dispatch(x, y, name)
   1204   is_tensor_y = isinstance(y, ops.Tensor)
   1205   if is_tensor_y:
-> 1206     return gen_math_ops.mul(x, y, name=name)
   1207   else:
   1208     assert isinstance(y, sparse_tensor.SparseTensor)  # Case: Dense * Sparse.

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_math_ops.py in mul(x, y, name)
   6696       else:
   6697         message = e.message
-> 6698       _six.raise_from(_core._status_to_exception(e.code, message), None)
   6699   # Add nodes to the TensorFlow graph.
   6700   _, _, _op = _op_def_lib._apply_op_helper(

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a int64 tensor but is a int32 tensor [Op:Mul] name: mul/
```

This can be fixed by cast ```num_tags``` in ```crf_binary_score```
```
num_tags = tf.cast(num_tags, dtype=tf.int64)
```

After fixing this error, another error comes up:
```
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-2-41e7d4eaa2ab> in <module>
      6     inputs=inputs,
      7     tag_indices=tags,
----> 8     sequence_lengths=seq_lens
      9 )

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_log_likelihood(inputs, tag_indices, sequence_lengths, transition_params)
    195     sequence_scores = crf_sequence_score(inputs, tag_indices, sequence_lengths,
    196                                          transition_params)
--> 197     log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)
    198 
    199     # Normalize the scores to get the log-likelihood per example.

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_log_norm(inputs, sequence_lengths, transition_params)
    161         return _single_seq_fn()
    162     else:
--> 163         return _multi_seq_fn()
    164 
    165 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in _multi_seq_fn()
    150 
    151         alphas = crf_forward(rest_of_input, first_input, transition_params,
--> 152                              sequence_lengths)
    153         log_norm = tf.reduce_logsumexp(alphas, [1])
    154         # Mask `log_norm` of the sequences with length <= zero.

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_addons/text/crf.py in crf_forward(inputs, state, transition_params, sequence_lengths)
    304     all_alphas = tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
    305     idxs = tf.stack(
--> 306         [tf.range(tf.shape(sequence_lengths)[0]), sequence_lengths], axis=1)
    307     return tf.gather_nd(all_alphas, idxs)
    308 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/util/dispatch.py in wrapper(*args, **kwargs)
    178     """Call target, and fall back on dispatchers if there is a TypeError."""
    179     try:
--> 180       return target(*args, **kwargs)
    181     except (TypeError, ValueError):
    182       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py in stack(values, axis, name)
   1163                        (axis, -expanded_num_dims, expanded_num_dims))
   1164 
-> 1165   return gen_array_ops.pack(values, axis=axis, name=name)
   1166 
   1167 

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_array_ops.py in pack(values, axis, name)
   6291       else:
   6292         message = e.message
-> 6293       _six.raise_from(_core._status_to_exception(e.code, message), None)
   6294   # Add nodes to the TensorFlow graph.
   6295   if not isinstance(values, (list, tuple)):

~/Deployment/tf_sequence_labeler/env2/lib/python3.6/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: cannot compute Pack as input #1(zero-based) was expected to be a int32 tensor but is a int64 tensor [Op:Pack] name: stack
```

This can be fixed by casting ```sequence_lengths``` to int32 in ```crf_forward``` function:
```
sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
```

Is this a bug, or I used it in a wrong way?


Discussion:
I think we should cast it internally to make data types the same.
cc @Squadrick @howl-anderson, who might be more familiar with CRF.

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
@@ -38,6 +38,8 @@ def crf_sequence_score(inputs, tag_indices, sequence_lengths,
     Returns:
       sequence_scores: A [batch_size] vector of unnormalized sequence scores.
     """
+    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
 
     # If max_seq_len is 1, we skip the score calculation and simply gather the
     # unary potentials of the single tag.
@@ -92,19 +94,19 @@ def crf_multitag_sequence_score(inputs, tag_bitmap, sequence_lengths,
     Returns:
       sequence_scores: A [batch_size] vector of unnormalized sequence scores.
     """
+    tag_bitmap = tf.cast(tag_bitmap, dtype=tf.bool)
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
+    filtered_inputs = tf.where(tag_bitmap, inputs,
+                               tf.fill(tf.shape(inputs), float("-inf")))
 
     # If max_seq_len is 1, we skip the score calculation and simply gather the
     # unary potentials of all active tags.
     def _single_seq_fn():
-        filtered_inputs = tf.where(tag_bitmap, inputs,
-                                   tf.fill(tf.shape(inputs), float("-inf")))
         return tf.reduce_logsumexp(
             filtered_inputs, axis=[1, 2], keepdims=False)
 
     def _multi_seq_fn():
         # Compute the logsumexp of all scores of sequences matching the given tags.
-        filtered_inputs = tf.where(tag_bitmap, inputs,
-                                   tf.fill(tf.shape(inputs), float("-inf")))
         return crf_log_norm(
             inputs=filtered_inputs,
             sequence_lengths=sequence_lengths,
@@ -127,6 +129,7 @@ def crf_log_norm(inputs, sequence_lengths, transition_params):
     Returns:
       log_norm: A [batch_size] vector of normalizers for a CRF.
     """
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
     # Split up the first and rest of the inputs in preparation for the forward
     # algorithm.
     first_input = tf.slice(inputs, [0, 0, 0], [-1, 1, -1])
@@ -183,10 +186,12 @@ def crf_log_likelihood(inputs,
       transition_params: A [num_tags, num_tags] transition matrix. This is
           either provided by the caller or created in this function.
     """
-    # Get shape information.
     num_tags = inputs.shape[2]
 
-    # Get the transition matrix if not provided.
+    # cast type to handle different types
+    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
+
     if transition_params is None:
         initializer = tf.keras.initializers.GlorotUniform()
         transition_params = tf.Variable(
@@ -211,6 +216,9 @@ def crf_unary_score(tag_indices, sequence_lengths, inputs):
     Returns:
       unary_scores: A [batch_size] vector of unary scores.
     """
+    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
+
     batch_size = tf.shape(inputs)[0]
     max_seq_len = tf.shape(inputs)[1]
     num_tags = tf.shape(inputs)[2]
@@ -245,7 +253,9 @@ def crf_binary_score(tag_indices, sequence_lengths, transition_params):
     Returns:
       binary_scores: A [batch_size] vector of binary scores.
     """
-    # Get shape information.
+    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
+
     num_tags = tf.shape(transition_params)[0]
     num_transitions = tf.shape(tag_indices)[1] - 1
 
@@ -288,6 +298,7 @@ def crf_forward(inputs, state, transition_params, sequence_lengths):
       new_alphas: A [batch_size, num_tags] matrix containing the
           new alpha values.
     """
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
 
     sequence_lengths = tf.maximum(
         tf.constant(0, dtype=sequence_lengths.dtype), sequence_lengths - 2)
@@ -399,6 +410,7 @@ def crf_decode_forward(inputs, state, transition_params, sequence_lengths):
       backpointers: A [batch_size, num_tags] matrix of backpointers.
       new_state: A [batch_size, num_tags] matrix of new score values.
     """
+    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
     mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
     crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params)
     crf_fwd_layer = tf.keras.layers.RNN(
@@ -446,6 +458,7 @@ def crf_decode(potentials, transition_params, sequence_length):
                   Contains the highest scoring tag indices.
       best_score: A [batch_size] vector, containing the score of `decode_tags`.
     """
+    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.


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
@@ -209,8 +209,11 @@ def testCrfLogLikelihood(self):
         transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]],
                                      dtype=np.float32)
         sequence_lengths = np.array(3, dtype=np.int32)
+        # TODO: https://github.com/PyCQA/pylint/issues/3139
+        # pylint: disable=E1136
         num_words = inputs.shape[0]
         num_tags = inputs.shape[1]
+        # pylint: enable=E1136
         all_sequence_log_likelihoods = []
 
         # Make sure all probabilities sum to 1.
@@ -241,8 +244,11 @@ def testViterbiDecode(self):
         transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]],
                                      dtype=np.float32)
         sequence_lengths = np.array(3, dtype=np.int32)
+        # TODO: https://github.com/PyCQA/pylint/issues/3139
+        # pylint: disable=E1136
         num_words = inputs.shape[0]
         num_tags = inputs.shape[1]
+        # pylint: enable=E1136
 
         all_sequence_scores = []
         all_sequences = []
@@ -347,6 +353,16 @@ def testCrfDecodeZeroSeqLength(self):
         self.assertEqual(len(tf_tags.shape), 2)
         self.assertEqual(len(tf_scores.shape), 1)
 
+    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
+
+        loss, _ = text.crf_log_likelihood(
+            inputs=inputs, tag_indices=tags, sequence_lengths=seq_lens)
+
 
 if __name__ == "__main__":
     tf.test.main()