Repository: tensorflow/addons

Issue:
Error with mixed precision
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.3 LTS
- TensorFlow version and how it was installed (source or binary):  https://github.com/horovod/horovod/blob/master/Dockerfile.gpu (changed TENSORFLOW_VERSION=2.2.0 and NCCL_VERSION=2.5.6-1+cuda10.1)
- TensorFlow-Addons version and how it was installed (source or binary): nightly (0.11.0-dev)
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

Code doesn't work when applying mixed precision.
Am I doing something wrong?

**Code to reproduce the issue**

```
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras.mixed_precision import experimental as mixed_precision

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

embedding_layer = tf.keras.layers.Embedding(37,
                                            37,
                                            name='embedding')

rnn_cell = tf.keras.layers.LSTMCell(10)
attention_mechanism = tfa.seq2seq.BahdanauAttention(units=10, name='BahdanauAtt')
attention_wrapper = tfa.seq2seq.AttentionWrapper(rnn_cell,
                                                 attention_mechanism,
                                                 attention_layer_size=10,
                                                 name='Att_wrapper')
mlp = tf.keras.layers.Dense(14, name='output_mlp')
decoder_cell = tfa.seq2seq.BasicDecoder(attention_wrapper,
                                        output_layer=mlp,
                                        sampler=tfa.seq2seq.ScheduledEmbeddingTrainingSampler(0.1,
                                                                                              embedding_layer),
                                        name='train_decoder')

@tf.function
def step(inputs, encoder_outputs):
    
    attention_mechanism.setup_memory(encoder_outputs)
    decoder_initial_state = attention_wrapper.get_initial_state(batch_size=32,
                                                                dtype=encoder_outputs.dtype)

    embeddings = embedding_layer(inputs)

    outputs = decoder_cell(embeddings,
                       initial_state=decoder_initial_state,
                       training=True)
    return outputs

encoder_outputs = tf.random.uniform([32, 17, 256], dtype=tf.float16)
inputs = tf.random.uniform([32, 27], maxval=37, dtype=tf.int32)
step(inputs, encoder_outputs)


```

**Other info / logs**

```
/usr/local/lib/python3.6/dist-packages/tensorflow_addons/seq2seq/decoder.py:163 call  *
        self,
    /usr/local/lib/python3.6/dist-packages/typeguard/__init__.py:262 wrapper  *
        retval = func(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow_addons/seq2seq/decoder.py:397 body  *
        (next_outputs, decoder_state, next_inputs, decoder_finished) = decoder.step(
    /usr/local/lib/python3.6/dist-packages/tensorflow_addons/seq2seq/basic_decoder.py:134 step  *
        cell_outputs, cell_state = self.cell(inputs, state, training=training)
    /usr/local/lib/python3.6/dist-packages/tensorflow_addons/seq2seq/attention_wrapper.py:2052 call  *
        attention, alignments, next_attention_state = self._attention_fn(
    /usr/local/lib/python3.6/dist-packages/tensorflow_addons/seq2seq/attention_wrapper.py:1595 _compute_attention  *
        context_ = tf.matmul(expanded_alignments, attention_mechanism.values)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py:180 wrapper  **
        return target(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py:2946 matmul
        a, b, adj_x=adjoint_a, adj_y=adjoint_b, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_math_ops.py:1526 batch_mat_mul_v2
        "BatchMatMulV2", x=x, y=y, adj_x=adj_x, adj_y=adj_y, name=name)
    /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 'BatchMatMulV2' Op has type float16 that does not match type float32 of argument 'x'.
```



Discussion:
The first issue is that we are forcing `float32` at https://github.com/tensorflow/addons/blob/master/tensorflow_addons/seq2seq/attention_wrapper.py#L559-L560.

Then we have a similar issue at https://github.com/tensorflow/addons/blob/master/tensorflow_addons/seq2seq/sampler.py#L394-L400
Yes, we should fix this. The seq2seq layers were designed before the new dtype policy.

While working on a fix, I find it a bit unexpected that `tf.keras.layers.Embedding` does not return `float16` here.

Resolving patch:
diff --git a/tensorflow_addons/conftest.py b/tensorflow_addons/conftest.py
--- a/tensorflow_addons/conftest.py
+++ b/tensorflow_addons/conftest.py
@@ -1,6 +1,7 @@
 from tensorflow_addons.utils.test_utils import (  # noqa: F401
     maybe_run_functions_eagerly,
     run_custom_and_py_ops,
+    run_with_mixed_precision_policy,
     pytest_make_parametrize_id,
     data_format,
     set_seeds,
diff --git a/tensorflow_addons/seq2seq/attention_wrapper.py b/tensorflow_addons/seq2seq/attention_wrapper.py
--- a/tensorflow_addons/seq2seq/attention_wrapper.py
+++ b/tensorflow_addons/seq2seq/attention_wrapper.py
@@ -110,8 +110,6 @@ def __init__(
         """
         self.query_layer = query_layer
         self.memory_layer = memory_layer
-        if self.memory_layer is not None and "dtype" not in kwargs:
-            kwargs["dtype"] = self.memory_layer.dtype
         super().__init__(**kwargs)
         self.default_probability_fn = probability_fn
         self.probability_fn = probability_fn
@@ -556,8 +554,6 @@ def __init__(
         def wrapped_probability_fn(score, _):
             return probability_fn(score)
 
-        if dtype is None:
-            dtype = tf.float32
         memory_layer = kwargs.pop("memory_layer", None)
         if not memory_layer:
             memory_layer = tf.keras.layers.Dense(
@@ -735,8 +731,6 @@ def __init__(
         def wrapped_probability_fn(score, _):
             return probability_fn(score)
 
-        if dtype is None:
-            dtype = tf.float32
         query_layer = kwargs.pop("query_layer", None)
         if not query_layer:
             query_layer = tf.keras.layers.Dense(
@@ -1093,8 +1087,6 @@ def __init__(
             creation.
         """
         # Set up the monotonic probability fn with supplied parameters
-        if dtype is None:
-            dtype = tf.float32
         wrapped_probability_fn = functools.partial(
             _monotonic_probability_fn,
             sigmoid_noise=sigmoid_noise,
@@ -1274,8 +1266,6 @@ def __init__(
             creation.
         """
         # Set up the monotonic probability fn with supplied parameters
-        if dtype is None:
-            dtype = tf.float32
         wrapped_probability_fn = functools.partial(
             _monotonic_probability_fn,
             sigmoid_noise=sigmoid_noise,
@@ -1729,12 +1719,13 @@ def cell_input_fn(inputs, attention):
                     "one integer per attention_mechanism, saw: %d vs %d"
                     % (len(attention_layer_sizes), len(attention_mechanisms))
                 )
+            dtype = kwargs.get("dtype", None)
             self._attention_layers = list(
                 tf.keras.layers.Dense(
                     attention_layer_size,
                     name="attention_layer",
                     use_bias=False,
-                    dtype=attention_mechanisms[i].dtype,
+                    dtype=dtype,
                 )
                 for i, attention_layer_size in enumerate(attention_layer_sizes)
             )
diff --git a/tensorflow_addons/seq2seq/sampler.py b/tensorflow_addons/seq2seq/sampler.py
--- a/tensorflow_addons/seq2seq/sampler.py
+++ b/tensorflow_addons/seq2seq/sampler.py
@@ -390,6 +390,9 @@ def maybe_sample():
             sample_ids_sampling = tf.gather_nd(sample_ids, where_sampling)
             inputs_not_sampling = tf.gather_nd(base_next_inputs, where_not_sampling)
             sampled_next_inputs = self.embedding_fn(sample_ids_sampling)
+            sampled_next_inputs = tf.cast(
+                sampled_next_inputs, inputs_not_sampling.dtype
+            )
             base_shape = tf.shape(base_next_inputs)
             return tf.scatter_nd(
                 indices=where_sampling, updates=sampled_next_inputs, shape=base_shape


Test patch:
diff --git a/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py b/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
--- a/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
+++ b/tensorflow_addons/seq2seq/tests/attention_wrapper_test.py
@@ -398,9 +398,12 @@ def _test_with_attention(
                 seed=1337
             )
 
+    policy = tf.keras.mixed_precision.experimental.global_policy()
     sampler = sampler_py.TrainingSampler()
     my_decoder = basic_decoder.BasicDecoder(cell=cell, sampler=sampler)
-    initial_state = cell.get_initial_state(dtype=tf.float32, batch_size=batch_size)
+    initial_state = cell.get_initial_state(
+        batch_size=batch_size, dtype=policy.compute_dtype
+    )
     final_outputs, final_state, _ = my_decoder(
         decoder_inputs,
         initial_state=initial_state,
@@ -427,7 +430,7 @@ def _test_with_attention(
         )
         tf.nest.assert_same_structure(
             cell.state_size,
-            cell.get_initial_state(batch_size=batch_size, dtype=tf.float32),
+            cell.get_initial_state(batch_size=batch_size, dtype=policy.compute_dtype),
         )
         # Remove the history from final_state for purposes of the
         # remainder of the tests.
@@ -535,32 +538,34 @@ def set_random_state_for_tf_and_np():
     DummyData2()
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_bahdanau_not_normalized():
     set_random_state_for_tf_and_np()
+    policy = tf.keras.mixed_precision.experimental.global_policy()
     create_attention_mechanism = wrapper.BahdanauAttention
     create_attention_kwargs = {"kernel_initializer": "ones"}
     expected_final_output = basic_decoder.BasicDecoderOutput(
         rnn_output=ResultSummary(
-            shape=(5, 3, 6), dtype=np.dtype(np.float32), mean=-0.003204414
+            shape=(5, 3, 6), dtype=policy.compute_dtype, mean=-0.003204414
         ),
         sample_id=ResultSummary(shape=(5, 3), dtype=np.dtype(np.int32), mean=3.2),
     )
     expected_final_state = wrapper.AttentionWrapperState(
         cell_state=[
-            ResultSummary(shape=(5, 9), dtype=np.dtype(np.float32), mean=0.40868404),
-            ResultSummary(shape=(5, 9), dtype=np.dtype(np.float32), mean=0.89017969),
+            ResultSummary(shape=(5, 9), dtype=policy.compute_dtype, mean=0.40868404),
+            ResultSummary(shape=(5, 9), dtype=policy.compute_dtype, mean=0.89017969),
         ],
         attention=ResultSummary(
-            shape=(5, 6), dtype=np.dtype(np.float32), mean=0.041453815
+            shape=(5, 6), dtype=policy.compute_dtype, mean=0.041453815
         ),
-        alignments=ResultSummary(shape=(5, 8), dtype=np.dtype(np.float32), mean=0.125),
+        alignments=ResultSummary(shape=(5, 8), dtype=policy.compute_dtype, mean=0.125),
         attention_state=ResultSummary(
-            shape=(5, 8), dtype=np.dtype(np.float32), mean=0.125
+            shape=(5, 8), dtype=policy.compute_dtype, mean=0.125
         ),
         alignment_history=(),
     )
     expected_final_alignment_history = ResultSummary(
-        shape=(3, 5, 8), dtype=np.dtype(np.float32), mean=0.125
+        shape=(3, 5, 8), dtype=policy.compute_dtype, mean=0.125
     )
 
     _test_with_attention(
@@ -609,27 +614,29 @@ def test_bahdanau_normalized():
     )
 
 
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_luong_not_normalized():
     set_random_state_for_tf_and_np()
+    policy = tf.keras.mixed_precision.experimental.global_policy()
     create_attention_mechanism = wrapper.LuongAttention
 
     expected_final_output = basic_decoder.BasicDecoderOutput(
         rnn_output=ResultSummary(
-            shape=(5, 3, 6), dtype=np.dtype("float32"), mean=-0.06124732
+            shape=(5, 3, 6), dtype=policy.compute_dtype, mean=-0.06124732
         ),
         sample_id=ResultSummary(shape=(5, 3), dtype=np.dtype("int32"), mean=2.73333333),
     )
     expected_final_state = wrapper.AttentionWrapperState(
         cell_state=[
-            ResultSummary(shape=(5, 9), dtype=np.dtype("float32"), mean=0.52021580),
-            ResultSummary(shape=(5, 9), dtype=np.dtype("float32"), mean=1.0964939),
+            ResultSummary(shape=(5, 9), dtype=policy.compute_dtype, mean=0.52021580),
+            ResultSummary(shape=(5, 9), dtype=policy.compute_dtype, mean=1.0964939),
         ],
         attention=ResultSummary(
-            shape=(5, 6), dtype=np.dtype("float32"), mean=-0.0318060
+            shape=(5, 6), dtype=policy.compute_dtype, mean=-0.0318060
         ),
-        alignments=ResultSummary(shape=(5, 8), dtype=np.dtype("float32"), mean=0.125),
+        alignments=ResultSummary(shape=(5, 8), dtype=policy.compute_dtype, mean=0.125),
         attention_state=ResultSummary(
-            shape=(5, 8), dtype=np.dtype("float32"), mean=0.125
+            shape=(5, 8), dtype=policy.compute_dtype, mean=0.125
         ),
         alignment_history=(),
     )
diff --git a/tensorflow_addons/seq2seq/tests/decoder_test.py b/tensorflow_addons/seq2seq/tests/decoder_test.py
--- a/tensorflow_addons/seq2seq/tests/decoder_test.py
+++ b/tensorflow_addons/seq2seq/tests/decoder_test.py
@@ -127,3 +127,31 @@ def test_dynamic_decode_rnn_with_training_helper_matches_dynamic_rnn(
     )
     if use_sequence_length:
         np.testing.assert_allclose(final_decoder_state, final_rnn_state)
+
+
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
+def test_dynamic_decode_rnn_with_scheduled_embedding_training_sampler():
+    policy = tf.keras.mixed_precision.experimental.global_policy()
+    sequence_length = [3, 4, 3, 1]
+    batch_size = 4
+    input_depth = 7
+    cell_depth = 10
+    vocab_size = 12
+    max_time = max(sequence_length)
+
+    embedding = tf.keras.layers.Embedding(vocab_size, input_depth)
+    cell = tf.keras.layers.LSTMCell(cell_depth)
+    sampler = sampler_py.ScheduledEmbeddingTrainingSampler(
+        sampling_probability=tf.constant(1.0), embedding_fn=embedding
+    )
+    my_decoder = basic_decoder.BasicDecoder(cell=cell, sampler=sampler)
+
+    inputs = tf.random.uniform([batch_size, max_time, input_depth])
+    initial_state = cell.get_initial_state(
+        batch_size=batch_size, dtype=policy.compute_dtype
+    )
+    final_outputs, _, _ = my_decoder(
+        inputs, initial_state=initial_state, sequence_length=sequence_length
+    )
+
+    assert final_outputs.rnn_output.dtype == policy.compute_dtype
diff --git a/tensorflow_addons/utils/test_utils.py b/tensorflow_addons/utils/test_utils.py
--- a/tensorflow_addons/utils/test_utils.py
+++ b/tensorflow_addons/utils/test_utils.py
@@ -21,6 +21,8 @@
 import pytest
 import tensorflow as tf
 
+from distutils.version import LooseVersion
+
 from tensorflow_addons import options
 from tensorflow_addons.utils import resource_loader
 
@@ -100,6 +102,15 @@ def _restore_py_ops_value():
     request.addfinalizer(_restore_py_ops_value)
 
 
+@pytest.fixture(scope="function", params=["float32", "mixed_float16"])
+def run_with_mixed_precision_policy(request):
+    if is_gpu_available() and LooseVersion(tf.__version__) <= "2.2.0":
+        pytest.xfail("See https://github.com/tensorflow/tensorflow/issues/39775")
+    tf.keras.mixed_precision.experimental.set_policy(request.param)
+    yield
+    tf.keras.mixed_precision.experimental.set_policy("float32")
+
+
 @pytest.fixture(scope="function", params=["channels_first", "channels_last"])
 def data_format(request):
     return request.param