tfa.text.crf_decode doesn't work after enabling mixed_precision

[solisxortuss]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Catalina 10.15.6
• TensorFlow version and how it was installed (source or binary): tensorflow==2.1 installed with pip
• TensorFlow-Addons version and how it was installed (source or binary): tensorflow-addons==0.9.1 installed with pip
• Python version: python 3.6.5
• Is GPU used? (yes/no): no

Describe the bug

The problem occurs when I'm trying to run tfa.text.crf_decode(...) after setting policy to mixed_float16. I'm getting the error tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a half tensor but is a float tensor [Op:AddV2] name: rnn_2/add/
The type of the first argument potentials is originally tf.float16, it doesn't work to pass it as is and results in the error described earlier. I'm trying to fix it and explicitly converting it to tf.float32 using tf.cast(...), the type of the second argument transition_params is tf.float32 and the type of the last argument sequence_length is tf.int64, it still doesn't help. Also, when I'm trying to convert sequence_length to tf.int32 the error appears again; playing with converting potentials and transition_params to tf.float64 doesn't work too. When mixed_precision option is disabled passing tf.float32, tf.float32 and tf.int64 tensor types to this method works ok with no errors.

Code to reproduce the issue

https://colab.research.google.com/drive/13Ft2U3cMhe1fcAqrO4Z6TZMSGsBduM-v?usp=sharing

Other info / logs

tfa.text.crf_log_likelihood(...) with mixed_precision option enabled works fine after casting all the variables to tf.float32 and tf.int32.

[WindQAQ]

@easwica Thanks for reporting! I believe this is resolved in f429133 so can you try pip install tfa-nightly to see if the issue still exists?
As per #2231, the mixed precision behavior is different from TensorFlow semantics. I'll submit a PR for it.

[solisxortuss]

@WindQAQ thank you, I've tried it and it works. However, looks like the result doesn't have the tf.float16 type, it is still tf.float32 which is a problem, but at least the main functionality works.

[WindQAQ]

Yep, there are some issues of TensorFlow mixed precision autocasting so that we cannot really benefit from it. Let's try to improve it.

https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/engine/base_layer.py#L904-L915