Repository: tensorflow/addons

Issue:
float16 support for text/crf
text/crf.py doesn't support float16 training or TF Lite converter.
According to the error information, the type of  masks In line 256 and line 294 is tf.float32:
```
masks = tf.sequence_mask(
        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=tf.float32
    )
```
which raises error when doling multiply op, like in line 259: 
```unary_scores = tf.reduce_sum(unary_scores * masks, 1)```


**Which API type would this fall under (layer, metric, optimizer, etc.)**
text
![image](https://user-images.githubusercontent.com/13343428/88151388-37387000-cc35-11ea-963f-68585a69f0e6.png)




Discussion:
Can you share a very minimal but runnable dummy example to reproduce this?
> Can you share a very minimal but runnable dummy example to reproduce this?
I have implemented the example below. The "tf2crf" is a small package that I wrote for using CRF layer with tf 2.0 keras. It just call tfa.text.crf_decode and tfa.text.crf_log_likelihood for computing CRF scores and loss.
```
from tensorflow.keras.layers import Input, Embedding, Bidirectional, GRU, Dense
from tensorflow.keras.models import Model
from tf2crf import CRF
from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

def test():
    inputs = Input(shape=(None,), dtype='int32')
    output = Embedding(100, 40, trainable=True, mask_zero=True)(inputs)
    output = Bidirectional(GRU(64, return_sequences=True))(output)
    output = Dense(9, activation=None)(output)
    crf = CRF()
    output = crf(output)
    model = Model(inputs, output)
    model.compile(loss=crf.loss, optimizer='adam', metrics=[crf.accuracy])

    x = [[5, 2, 3] * 3] * 10
    y = [[1, 2, 3] * 3] * 10

    model.fit(x=x, y=y, epochs=2, batch_size=2)
    model.save('model')


if __name__ == '__main__':
    test()
```
When run the code, I got this error:
```
TypeError: true_fn and false_fn arguments to tf.cond must have the same number, type, and overall structure of return values.
    
    true_fn output: Tensor("Identity:0", shape=(None, None, 9), dtype=float16)
    false_fn output: Tensor("Identity:0", shape=(None, None, 9), dtype=float32)
    
    Error details:
    Tensor("Identity:0", shape=(None, None, 9), dtype=float16) and Tensor("Identity:0", shape=(None, None, 9), dtype=float32) have different types
```
The error raised when called tfa.text.crf_decode in the CRF layer
/cc @Squadrick 
I think we can change these two dtypes from `tf.float32` to `unary_scores.dtype` and `binary_scores.dtype` respectively.
@xuxingya Can you try it or even issue an PR for it? Thanks!

https://github.com/tensorflow/addons/blob/master/tensorflow_addons/text/crf.py#L256
https://github.com/tensorflow/addons/blob/master/tensorflow_addons/text/crf.py#L294
@WindQAQ no we need something else.
I am facing the same issue when using the recently added CRF layer in the TFA nightly package from this commit:
https://github.com/tensorflow/addons/commit/4f65776d7ab6d31884aa91445fd9fe94a504bac1
Are there any plans to fixing this. I'm happy to help out! @bhack @WindQAQ 
> I am facing the same issue when using the recently added CRF layer in the TFA nightly package from this commit:
> [4f65776](https://github.com/tensorflow/addons/commit/4f65776d7ab6d31884aa91445fd9fe94a504bac1)
> Are there any plans to fixing this. I'm happy to help out! @bhack @WindQAQ

I'll investigate it later, could you provide runnable minimal code snippet to reproduce the issue? seems that the code above requires some customized model. Thanks.
> > I am facing the same issue when using the recently added CRF layer in the TFA nightly package from this commit:
> > [4f65776](https://github.com/tensorflow/addons/commit/4f65776d7ab6d31884aa91445fd9fe94a504bac1)
> > Are there any plans to fixing this. I'm happy to help out! @bhack @WindQAQ
> 
> I'll investigate it later, could you provide runnable minimal code snippet to reproduce the issue? seems that the code above requires some customized model. Thanks.

Here's an example code that uses the recently added CRF layer as well as the error produced.

```
from tensorflow.keras.layers import Input, Embedding, Bidirectional, LSTM, Dense
from tensorflow.keras.models import Model
from tensorflow_addons.layers.crf import CRF
from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

def test():
    inputs = Input(shape=(None,), dtype='int32')
    output = Embedding(100, 40, trainable=True, mask_zero=True)(inputs)
    output = Bidirectional(LSTM(10, return_sequences=True))(output)
    output = Dense(9, activation=None)(output)
    crf = CRF()
    output = crf(output)
    model = Model(inputs, output)
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics='acc')
    x = [[5, 2, 3] * 3] * 10
    y = [[1, 2, 3] * 3] * 10

    model.fit(x=x, y=y, epochs=2, batch_size=2)


if __name__ == '__main__':
    test()

```

Error:
```
...
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:506 _apply_op_helper
        inferred_from[input_arg.type_attr]))

    TypeError: Input 'y' of 'Mul' Op has type float64 that does not match type int32 of argument 'x'.
```
>     x = [[5, 2, 3] * 3] * 10
>     y = [[1, 2, 3] * 3] * 10

I'm not so familiar with CRF, but can it really be used to compute cross entropy with targets? Seems that it's output is sparse integer.

Resolving patch:
diff --git a/tensorflow_addons/layers/crf.py b/tensorflow_addons/layers/crf.py
--- a/tensorflow_addons/layers/crf.py
+++ b/tensorflow_addons/layers/crf.py
@@ -189,8 +189,8 @@ def expand_scalar_to_3d(x):
             # expand tensor from shape (x, ) to (1, 1, x)
             return tf.reshape(x, (1, 1, -1))
 
-        start = expand_scalar_to_3d(start)
-        end = expand_scalar_to_3d(end)
+        start = tf.cast(expand_scalar_to_3d(start), potentials.dtype)
+        end = tf.cast(expand_scalar_to_3d(end), potentials.dtype)
         if mask is None:
             potentials = tf.concat(
                 [potentials[:, :1, :] + start, potentials[:, 1:, :]], axis=1
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
@@ -108,7 +108,9 @@ def crf_multitag_sequence_score(
     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"))
+        tag_bitmap,
+        inputs,
+        tf.fill(tf.shape(inputs), tf.cast(float("-inf"), inputs.dtype)),
     )
 
     # If max_seq_len is 1, we skip the score calculation and simply gather the
@@ -198,18 +200,21 @@ def crf_log_likelihood(
       transition_params: A [num_tags, num_tags] transition matrix. This is
           either provided by the caller or created in this function.
     """
+    inputs = tf.convert_to_tensor(inputs)
+
     num_tags = inputs.shape[2]
 
     # cast type to handle different types
     tag_indices = tf.cast(tag_indices, dtype=tf.int32)
     sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
 
+    # TODO(windqaq): re-evaluate if `transition_params` can be `None`.
     if transition_params is None:
         initializer = tf.keras.initializers.GlorotUniform()
         transition_params = tf.Variable(
             initializer([num_tags, num_tags]), "transitions"
         )
-
+    transition_params = tf.cast(transition_params, inputs.dtype)
     sequence_scores = crf_sequence_score(
         inputs, tag_indices, sequence_lengths, transition_params
     )
@@ -253,7 +258,7 @@ def crf_unary_score(
     )
 
     masks = tf.sequence_mask(
-        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=tf.float32
+        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=unary_scores.dtype
     )
 
     unary_scores = tf.reduce_sum(unary_scores * masks, 1)
@@ -291,7 +296,7 @@ def crf_binary_score(
     binary_scores = tf.gather(flattened_transition_params, flattened_transition_indices)
 
     masks = tf.sequence_mask(
-        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=tf.float32
+        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=binary_scores.dtype
     )
     truncated_masks = tf.slice(masks, [0, 1], [-1, -1])
     binary_scores = tf.reduce_sum(binary_scores * truncated_masks, 1)
@@ -416,8 +421,8 @@ def call(self, inputs, state):
           new_state: A [batch_size, num_tags] matrix of new score values.
         """
         state = tf.expand_dims(state[0], 2)
-        transition_scores = state + self._transition_params
-        new_state = inputs + tf.reduce_max(transition_scores, [1])
+        transition_scores = state + tf.cast(self._transition_params, state.dtype)
+        new_state = tf.cast(inputs, state.dtype) + tf.reduce_max(transition_scores, [1])
         backpointers = tf.argmax(transition_scores, 1)
         backpointers = tf.cast(backpointers, dtype=tf.int32)
         return backpointers, new_state


Test patch:
diff --git a/tensorflow_addons/layers/tests/crf_test.py b/tensorflow_addons/layers/tests/crf_test.py
--- a/tensorflow_addons/layers/tests/crf_test.py
+++ b/tensorflow_addons/layers/tests/crf_test.py
@@ -25,6 +25,7 @@
 
 from tensorflow_addons.layers.crf import CRF
 from tensorflow_addons.text.crf import crf_log_likelihood
+from tensorflow_addons.utils import test_utils
 
 
 def get_test_data():
@@ -79,6 +80,7 @@ def get_test_data_extended():
     return logits, tags, transitions, boundary_values, crf_layer
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_keras_model_inference():
     logits, _, _, _, crf_layer = get_test_data_extended()
 
@@ -90,8 +92,8 @@ def test_keras_model_inference():
     model(logits).numpy()
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_unmasked_viterbi_decode():
-
     x_np, y_np = get_test_data()
 
     transitions = np.ones([5, 5])
@@ -162,6 +164,7 @@ def test_step(self, data):
         return {"crf_loss_val": crf_loss, "internal_losses_val": internal_losses}
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_traing():
     x_np, y_np = get_test_data()
     get_some_model(x_np, y_np)
@@ -182,6 +185,7 @@ def get_some_model(x_np, y_np, sanity_check=True):
     return model
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_mask_right_padding():
     x_np, y_np = get_test_data()
     mask = np.array([[1, 1, 1], [1, 1, 0]])
@@ -206,6 +210,7 @@ def test_mask_right_padding():
     model.predict(x_np)
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_mask_left_padding():
     x_np, y_np = get_test_data()
     mask = np.array([[0, 1, 1], [1, 1, 1]])
@@ -269,6 +274,7 @@ def test_serialization(inference_only):
         assert original_loss == clone_loss
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_numerical_accuracy():
     logits, tags, transitions, boundary_values, crf_layer = get_test_data_extended()
@@ -288,8 +294,8 @@ def test_numerical_accuracy():
     )
     unbatched_log_likelihood = -2 * log_likelihood
 
-    np.testing.assert_allclose(
-        expected_log_likelihood, unbatched_log_likelihood, rtol=2e-7
+    test_utils.assert_allclose_according_to_type(
+        expected_log_likelihood, unbatched_log_likelihood, rtol=5e-5
     )
 
 
diff --git a/tensorflow_addons/text/tests/crf_test.py b/tensorflow_addons/text/tests/crf_test.py
--- a/tensorflow_addons/text/tests/crf_test.py
+++ b/tensorflow_addons/text/tests/crf_test.py
@@ -16,10 +16,12 @@
 
 import itertools
 
+import pytest
 import numpy as np
 import tensorflow as tf
 
 from tensorflow_addons import text
+from tensorflow_addons.utils import test_utils
 
 
 def calculate_sequence_score(inputs, transition_params, tag_indices, sequence_lengths):
@@ -33,16 +35,17 @@ def calculate_sequence_score(inputs, transition_params, tag_indices, sequence_le
     return expected_unary_score + expected_binary_score
 
 
-def test_crf_sequence_score():
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_sequence_score(dtype):
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     # Test both the length-1 and regular cases.
     sequence_lengths_list = [
         np.array(3, dtype=np.int32),
         np.array(1, dtype=np.int32),
     ]
     inputs_list = [
-        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32),
-        np.array([[4, 5, -3]], dtype=np.float32),
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[4, 5, -3]], dtype=dtype),
     ]
     tag_indices_list = [
         np.array([1, 2, 1, 0], dtype=np.int32),
@@ -62,19 +65,22 @@ def test_crf_sequence_score():
         expected_sequence_score = calculate_sequence_score(
             inputs, transition_params, tag_indices, sequence_lengths
         )
-        np.testing.assert_allclose(sequence_score, expected_sequence_score)
+        test_utils.assert_allclose_according_to_type(
+            sequence_score, expected_sequence_score
+        )
 
 
-def test_crf_multi_tag_sequence_score():
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_multi_tag_sequence_score(dtype):
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     # Test both the length-1 and regular cases.
     sequence_lengths_list = [
         np.array(3, dtype=np.int32),
         np.array(1, dtype=np.int32),
     ]
     inputs_list = [
-        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32),
-        np.array([[4, 5, -3]], dtype=np.float32),
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[4, 5, -3]], dtype=dtype),
     ]
     tag_bitmap_list = [
         np.array(
@@ -111,11 +117,14 @@ def test_crf_multi_tag_sequence_score():
         expected_log_sum_exp_sequence_scores = np.logaddexp.reduce(
             expected_sequence_scores
         )
-        np.testing.assert_allclose(sequence_score, expected_log_sum_exp_sequence_scores)
+        test_utils.assert_allclose_according_to_type(
+            sequence_score, expected_log_sum_exp_sequence_scores
+        )
 
 
-def test_crf_unary_score():
-    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_unary_score(dtype):
+    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype)
     for dtype in (np.int32, np.int64):
         tag_indices = np.array([1, 2, 1, 0], dtype=dtype)
         sequence_lengths = np.array(3, dtype=np.int32)
@@ -128,12 +137,13 @@ def test_crf_unary_score():
         expected_unary_score = sum(
             inputs[i][tag_indices[i]] for i in range(sequence_lengths)
         )
-        np.testing.assert_allclose(unary_score, expected_unary_score)
+        test_utils.assert_allclose_according_to_type(unary_score, expected_unary_score)
 
 
-def test_crf_binary_score():
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_binary_score(dtype):
     tag_indices = np.array([1, 2, 1, 0], dtype=np.int32)
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     sequence_lengths = np.array(3, dtype=np.int32)
     binary_score = text.crf_binary_score(
         tag_indices=tf.expand_dims(tag_indices, 0),
@@ -145,19 +155,20 @@ def test_crf_binary_score():
         transition_params[tag_indices[i], tag_indices[i + 1]]
         for i in range(sequence_lengths - 1)
     )
-    np.testing.assert_allclose(binary_score, expected_binary_score)
+    test_utils.assert_allclose_according_to_type(binary_score, expected_binary_score)
 
 
-def test_crf_log_norm():
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_log_norm(dtype):
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     # Test both the length-1 and regular cases.
     sequence_lengths_list = [
         np.array(3, dtype=np.int32),
         np.array(1, dtype=np.int64),
     ]
     inputs_list = [
-        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32),
-        np.array([[3, -1, 3]], dtype=np.float32),
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[3, -1, 3]], dtype=dtype),
     ]
     tag_indices_list = [
         np.array([1, 2, 1, 0], dtype=np.int32),
@@ -192,23 +203,25 @@ def test_crf_log_norm():
         )
         log_norm = tf.squeeze(log_norm, [0])
 
-        np.testing.assert_allclose(log_norm, brute_force_log_norm)
+        test_utils.assert_allclose_according_to_type(log_norm, brute_force_log_norm)
 
 
-def test_crf_log_norm_zero_seq_length():
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_log_norm_zero_seq_length(dtype):
     """Test `crf_log_norm` when `sequence_lengths` contains one or more
     zeros."""
-    inputs = tf.constant(np.ones([2, 10, 5], dtype=np.float32))
-    transition_params = tf.constant(np.ones([5, 5], dtype=np.float32))
+    inputs = tf.constant(np.ones([2, 10, 5], dtype=dtype))
+    transition_params = tf.constant(np.ones([5, 5], dtype=dtype))
     sequence_lengths = tf.constant(np.zeros([2], dtype=np.int32))
-    expected_log_norm = np.zeros([2], dtype=np.float32)
+    expected_log_norm = np.zeros([2], dtype=dtype)
     log_norm = text.crf_log_norm(inputs, sequence_lengths, transition_params)
-    np.testing.assert_allclose(log_norm, expected_log_norm)
+    test_utils.assert_allclose_according_to_type(log_norm, expected_log_norm)
 
 
-def test_crf_log_likelihood():
-    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32)
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_log_likelihood(dtype):
+    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype)
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     sequence_lengths = np.array(3, dtype=np.int32)
 
     num_words = inputs.shape[0]
@@ -227,7 +240,9 @@ def test_crf_log_likelihood():
         )
         all_sequence_log_likelihoods.append(sequence_log_likelihood)
     total_log_likelihood = tf.reduce_logsumexp(all_sequence_log_likelihoods)
-    np.testing.assert_allclose(total_log_likelihood, 0.0, 1e-6, 1e-6)
+    test_utils.assert_allclose_according_to_type(
+        total_log_likelihood, 0.0, rtol=1e-6, atol=1e-6, half_rtol=2e-3, half_atol=2e-3
+    )
 
     # check if `transition_params = None` raises an error
     text.crf_log_likelihood(
@@ -237,9 +252,10 @@ def test_crf_log_likelihood():
     )
 
 
-def test_viterbi_decode():
-    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32)
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_viterbi_decode(dtype):
+    inputs = np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype)
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     sequence_lengths = np.array(3, dtype=np.int32)
     num_words = inputs.shape[0]
     num_tags = inputs.shape[1]
@@ -269,20 +285,21 @@ def test_viterbi_decode():
         inputs[:sequence_lengths], transition_params
     )
 
-    np.testing.assert_allclose(actual_max_score, expected_max_score)
+    test_utils.assert_allclose_according_to_type(actual_max_score, expected_max_score)
     assert actual_max_sequence == expected_max_sequence[:sequence_lengths]
 
 
-def test_crf_decode():
-    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=np.float32)
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+def test_crf_decode(dtype):
+    transition_params = np.array([[-3, 5, -2], [3, 4, 1], [1, 2, 1]], dtype=dtype)
     # Test both the length-1 and regular cases.
     sequence_lengths_list = [
         np.array(3, dtype=np.int32),
         np.array(1, dtype=np.int64),
     ]
     inputs_list = [
-        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32),
-        np.array([[-1, 2, 1]], dtype=np.float32),
+        np.array([[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=dtype),
+        np.array([[-1, 2, 1]], dtype=dtype),
     ]
     tag_indices_list = [
         np.array([1, 2, 1, 0], dtype=np.int32),
@@ -324,7 +341,9 @@ def test_crf_decode():
         actual_max_sequence = tf.squeeze(actual_max_sequence, [0])
         actual_max_score = tf.squeeze(actual_max_score, [0])
 
-        np.testing.assert_allclose(actual_max_score, expected_max_score, 1e-6, 1e-6)
+        test_utils.assert_allclose_according_to_type(
+            actual_max_score, expected_max_score, 1e-6, 1e-6
+        )
         assert (
             list(actual_max_sequence[:sequence_lengths])
             == expected_max_sequence[:sequence_lengths]