Repository: tensorflow/addons

Issue:
Significant LazyAdam optimizer performance degradation since PR#1988
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Reproducible on Colab
- TensorFlow version and how it was installed (source or binary): TF 2.3.0
- TensorFlow-Addons version and how it was installed (source or binary): TF 0.11.2
- Python version: 3.6.9
- Is GPU used? (yes/no): no (but issue observed on GPU as well)

**Describe the bug**

PR [#1988](https://github.com/tensorflow/addons/pull/1988/files) replaces calls to resource scatter update/sub/add from `tf.raw_ops` with calls to similar methods from ancestor class [OptimizerV2](https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L1149-L1157).
These differ in that the OptimizerV2 method calls `.value()` on the input resource and returns a `Tensor`, whereas the `raw_ops` method returns an `Operation`.

The result is a major performance penalty with both CPU and GPU runtimes (in fact when using a GPU I have observed GPU utilization to drop to near 0%).

**Code to reproduce the issue**

Issue reproduced in this Colab on a CPU runtime with both with the Keras API and the Estimator API:
https://colab.research.google.com/drive/1IxPrQiGQn9Wgn9MtMhVTh0rdLlgZkMYo?usp=sharing


Discussion:
google-cla[bot]: 
Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

:memo: **Please visit <https://cla.developers.google.com/> to sign.**

Once you've signed (or fixed any issues), please reply here with `@googlebot I signed it!` and we'll verify it.

----

#### What to do if you already signed the CLA

##### Individual signers

*   It's possible we don't have your GitHub username or you're using a different email address on your commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).

##### Corporate signers

*   Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to [go/cla#troubleshoot](http://go/cla#troubleshoot) ([Public version](https://opensource.google/docs/cla/#troubleshoot)).
*   The email used to register you as an authorized contributor must be the email used for the Git commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).
*   The email used to register you as an authorized contributor must also be [attached to your GitHub account](https://github.com/settings/emails).
		

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2Ftensorflow%2Faddons%2Fpull%2F2274) for more info**.

<!-- need_sender_cla -->
bot-of-gabrieldemarmiesse: @ssaishruthi 

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.

edend10: @googlebot I signed it!

Patch:
diff --git a/tensorflow_addons/callbacks/average_model_checkpoint.py b/tensorflow_addons/callbacks/average_model_checkpoint.py
index aef5ffb..5c55c76 100644
--- a/tensorflow_addons/callbacks/average_model_checkpoint.py
+++ b/tensorflow_addons/callbacks/average_model_checkpoint.py
@@ -61,25 +61,35 @@ class AverageModelCheckpoint(tf.keras.callbacks.ModelCheckpoint):
             **kwargs,
         )
 
+    def _get_optimizer(self):
+        optimizer = self.model.optimizer
+        # TODO(fsx950223): change _optimizer to inner_optimizer after tf2.3 is not support
+        if type(optimizer).__name__ == "LossScaleOptimizer":
+            optimizer = optimizer._optimizer
+
+        return optimizer
+
     def set_model(self, model):
-        if not isinstance(model.optimizer, AveragedOptimizerWrapper):
+        super().set_model(model)
+        optimizer = self._get_optimizer()
+        if not isinstance(optimizer, AveragedOptimizerWrapper):
             raise TypeError(
                 "AverageModelCheckpoint is only used when training"
                 "with MovingAverage or StochasticAverage"
             )
-        return super().set_model(model)
 
     def _save_model(self, epoch, logs):
-        assert isinstance(self.model.optimizer, AveragedOptimizerWrapper)
+        optimizer = self._get_optimizer()
+        assert isinstance(optimizer, AveragedOptimizerWrapper)
 
         if self.update_weights:
-            self.model.optimizer.assign_average_vars(self.model.variables)
+            optimizer.assign_average_vars(self.model.variables)
             return super()._save_model(epoch, logs)
         else:
             # Note: `model.get_weights()` gives us the weights (non-ref)
             # whereas `model.variables` returns references to the variables.
             non_avg_weights = self.model.get_weights()
-            self.model.optimizer.assign_average_vars(self.model.variables)
+            optimizer.assign_average_vars(self.model.variables)
             # result is currently None, since `super._save_model` doesn't
             # return anything, but this may change in the future.
             result = super()._save_model(epoch, logs)
diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py
index 50b6152..495e02b 100644
--- a/tensorflow_addons/optimizers/average_wrapper.py
+++ b/tensorflow_addons/optimizers/average_wrapper.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import abc
+import warnings
 
 import tensorflow as tf
 from tensorflow_addons.utils import types
@@ -103,14 +104,17 @@ class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCM
         model.save('model.h5')
         ```
         """
-        assign_op = tf.group(
-            [
-                var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
-                for var in var_list
-                if var.trainable
-            ]
-        )
-        return assign_op
+        assign_ops = []
+        for var in var_list:
+            try:
+                assign_ops.append(
+                    var.assign(
+                        self.get_slot(var, "average"), use_locking=self._use_locking
+                    )
+                )
+            except Exception as e:
+                warnings.warn("Unable to assign average slot to {} : {}".format(var, e))
+        return tf.group(assign_ops)
 
     def get_config(self):
         config = {
diff --git a/tensorflow_addons/optimizers/lazy_adam.py b/tensorflow_addons/optimizers/lazy_adam.py
index 96a37d2..7b4cf8e 100644
--- a/tensorflow_addons/optimizers/lazy_adam.py
+++ b/tensorflow_addons/optimizers/lazy_adam.py
@@ -119,7 +119,26 @@ class LazyAdam(tf.keras.optimizers.Adam):
         v_update_op = self._resource_scatter_update(v, indices, v_t_slice)
 
         # \\(variable += -learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
-        var_slice = -lr * m_t_slice / (tf.math.sqrt(v_t_slice) + epsilon_t)
-        var_update_op = self._resource_scatter_add(var, indices, var_slice)
+        var_slice = lr * m_t_slice / (tf.math.sqrt(v_t_slice) + epsilon_t)
+        var_update_op = self._resource_scatter_sub(var, indices, var_slice)
 
         return tf.group(*[var_update_op, m_update_op, v_update_op])
+
+    def _resource_scatter_update(self, resource, indices, update):
+        return self._resource_scatter_operate(
+            resource, indices, update, tf.raw_ops.ResourceScatterUpdate
+        )
+
+    def _resource_scatter_sub(self, resource, indices, update):
+        return self._resource_scatter_operate(
+            resource, indices, update, tf.raw_ops.ResourceScatterSub
+        )
+
+    def _resource_scatter_operate(self, resource, indices, update, resource_scatter_op):
+        resource_update_kwargs = {
+            "resource": resource.handle,
+            "indices": indices,
+            "updates": update,
+        }
+
+        return resource_scatter_op(**resource_update_kwargs)
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
index 8bc9901..63f0e6a 100644
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -443,8 +443,10 @@ 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 + tf.cast(self._transition_params, state.dtype)
-        new_state = tf.cast(inputs, state.dtype) + tf.reduce_max(transition_scores, [1])
+        transition_scores = state + tf.cast(
+            self._transition_params, self._compute_dtype
+        )
+        new_state = inputs + tf.reduce_max(transition_scores, [1])
         backpointers = tf.argmax(transition_scores, 1)
         backpointers = tf.cast(backpointers, dtype=tf.int32)
         return backpointers, new_state
@@ -485,9 +487,9 @@ def crf_decode_forward(
     """
     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_cell = CrfDecodeForwardRnnCell(transition_params, dtype=inputs.dtype)
     crf_fwd_layer = tf.keras.layers.RNN(
-        crf_fwd_cell, return_sequences=True, return_state=True
+        crf_fwd_cell, return_sequences=True, return_state=True, dtype=inputs.dtype
     )
     return crf_fwd_layer(inputs, state, mask=mask)
 


Test patch:
diff --git a/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py b/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
index 4dbcf2c..8cbe970 100644
--- a/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
+++ b/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
@@ -195,3 +195,25 @@ def test_invalid_save_freq(tmp_path):
         AverageModelCheckpoint(
             update_weights=True, filepath=test_model_filepath, save_freq=save_freq
         )
+
+
+def test_loss_scale_optimizer(tmp_path):
+    test_model_filepath = str(tmp_path / "test_model.{epoch:02d}.h5")
+    moving_avg = MovingAverage(tf.keras.optimizers.SGD(lr=2.0), average_decay=0.5)
+    optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
+        moving_avg, "dynamic"
+    )
+    x, y, model = get_data_and_model(optimizer)
+    save_freq = "epoch"
+    avg_model_ckpt = AverageModelCheckpoint(
+        update_weights=False, filepath=test_model_filepath, save_freq=save_freq
+    )
+    model.fit(
+        x,
+        y,
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        validation_data=(x, y),
+        callbacks=[avg_model_ckpt],
+    )
+    assert not os.path.exists(test_model_filepath)
diff --git a/tensorflow_addons/optimizers/tests/moving_average_test.py b/tensorflow_addons/optimizers/tests/moving_average_test.py
index bf41616..bddce8b 100644
--- a/tensorflow_addons/optimizers/tests/moving_average_test.py
+++ b/tensorflow_addons/optimizers/tests/moving_average_test.py
@@ -295,7 +295,6 @@ def test_swap_weights(device):
 
 @pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_model_mixed_precision():
-    tf.keras.mixed_precision.experimental.set_policy("mixed_float16")
     x = np.random.standard_normal((10000, 3))
     w = np.random.standard_normal((3, 1))
     y = np.dot(x, w) + np.random.standard_normal((10000, 1)) * 1e-4
@@ -303,3 +302,35 @@ def test_model_mixed_precision():
     model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
     model.compile(MovingAverage("sgd"), loss="mse")
     model.fit(x, y, epochs=3)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_no_average_slot():
+    max_features = 5000
+    max_len = 4
+    embedding_dims = 2
+
+    # Some preprocessing layers have TrackableWeightHandler.
+    # They are returned when using model.variables
+    # but it's unable to assign average slot to them.
+    vectorize_layer = tf.keras.layers.experimental.preprocessing.TextVectorization(
+        max_tokens=max_features, output_mode="int", output_sequence_length=max_len
+    )
+
+    vectorize_layer.adapt(["foo", "bar", "baz"])
+
+    model = tf.keras.models.Sequential(
+        [
+            tf.keras.Input(shape=(1,), dtype=tf.string),
+            vectorize_layer,
+            tf.keras.layers.Embedding(max_features + 1, embedding_dims),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+
+    optimizer = MovingAverage("sgd")
+
+    model.compile(optimizer, loss="mse")
+    model.fit(x=["foo", "bar", "baz"], y=[0.0, 1.0, 2.0], epochs=1)
+
+    optimizer.assign_average_vars(model.variables)