Merge pull request #17225 from lgeiger:fix-mixed-precision-ema (#17226)

PiperOrigin-RevId: 486999108

Co-authored-by: TensorFlower Gardener <gardener@tensorflow.org>

Author: qlzh727

keras/engine/training_test.py

@@ -2172,9 +2172,13 @@ def metrics(self):
         )
 
     @test_combinations.run_all_keras_modes(always_skip_v1=True)
-    def test_ema_overwrite(self):
+    @parameterized.named_parameters(
+        ("mixed_float16", "mixed_float16"), ("float32", "float32")
+    )
+    def test_ema_overwrite(self, test_policy):
         if not tf.__internal__.tf2.enabled():
             self.skipTest("EMA optimizer is only available in TF2.")
+        policy.set_global_policy(test_policy)
         model = sequential.Sequential()
         model.add(input_layer.Input(shape=(4,)))
         model.add(layers_module.Dense(1, activation="relu"))
@@ -2188,6 +2192,7 @@ def test_ema_overwrite(self):
         history = model.fit(dataset, epochs=2, steps_per_epoch=10)
         self.assertLen(history.history["loss"], 2)
         self.assertAllClose(initial_value, model.trainable_variables[0])
+        policy.set_global_policy("float32")
 
     @test_combinations.run_all_keras_modes(always_skip_v1=True)
     def test_get_verbosity(self):

keras/mixed_precision/loss_scale_optimizer.py

@@ -1415,6 +1415,9 @@ def ema_momentum(self):
     def ema_momentum(self, ema_momentum):
         self._optimizer.ema_momentum = ema_momentum
 
+    def finalize_variable_values(self, var_list):
+        self._optimizer.finalize_variable_values(var_list)
+
 
 class FakeOptimizerForRestoration(tf.__internal__.tracking.Trackable):
     """A fake optimizer used to support restoring TensorFlow 2.2 checkpoints.