Add experimental_aggregate_gradients support (#2137)

* Add experimental_aggregate_gradients support

* format code

* fix test case

Author: fsx950223

tensorflow_addons/optimizers/average_wrapper.py

@@ -65,9 +65,9 @@ def _create_hypers(self):
     def _prepare(self, var_list):
         return self._optimizer._prepare(var_list=var_list)
 
-    def apply_gradients(self, grads_and_vars, name=None):
+    def apply_gradients(self, grads_and_vars, name=None, **kwargs):
         self._optimizer._iterations = self.iterations
-        return super().apply_gradients(grads_and_vars, name)
+        return super().apply_gradients(grads_and_vars, name, **kwargs)
 
     @abc.abstractmethod
     def average_op(self, var, average_var):

tensorflow_addons/optimizers/tests/moving_average_test.py

@@ -259,3 +259,15 @@ def apply_gradients():
 
     np.testing.assert_allclose(var.read_value(), [0.8, 1.8])
     np.testing.assert_allclose(ema_var.read_value(), [0.9, 1.9])
+
+
+@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
+    model = tf.keras.Sequential()
+    model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
+    model.compile(MovingAverage("sgd"), loss="mse")
+    model.fit(x, y, epochs=3)