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:

Patch:
diff --git a/tensorflow_addons/layers/crf.py b/tensorflow_addons/layers/crf.py
index 998d56b..7449ce7 100644
--- a/tensorflow_addons/layers/crf.py
+++ b/tensorflow_addons/layers/crf.py
@@ -189,8 +189,8 @@ class CRF(tf.keras.layers.Layer):
             # 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
index afd5cc8..2707895 100644
--- 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 @@ class CrfDecodeForwardRnnCell(tf.keras.layers.AbstractRNNCell):
           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
index 7896854..aa0bb4f 100644
--- a/tensorflow_addons/layers/tests/crf_test.py
+++ b/tensorflow_addons/layers/tests/crf_test.py
@@ -25,6 +25,7 @@ import tensorflow as tf
 
 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 @@ class ModelWithCRFLoss(tf.keras.Model):
         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
index f2c75a8..fc92a54 100644
--- 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]