Code changes to get ready for an incoming Keras optimizer migration.

PiperOrigin-RevId: 476479516

official/core/base_trainer_test.py

@@ -319,7 +319,9 @@ class TrainerTest(tf.test.TestCase, parameterized.TestCase):
         self.assertFalse(trainer.optimizer.dynamic)
         self.assertEqual(trainer.optimizer.initial_scale, loss_scale)
     else:
-      self.assertIsInstance(trainer.optimizer, tf.keras.optimizers.SGD)
+      self.assertIsInstance(
+          trainer.optimizer,
+          (tf.keras.optimizers.SGD, tf.keras.optimizers.legacy.SGD))
 
     metrics = trainer.train(tf.convert_to_tensor(5, dtype=tf.int32))
     self.assertIn('training_loss', metrics)

official/modeling/fast_training/progressive/trainer_test.py

@@ -226,9 +226,13 @@ class TrainerWithMaskedLMTaskTest(tf.test.TestCase, parameterized.TestCase):
     task = TestPolicy(None, config.task)
     trainer = trainer_lib.ProgressiveTrainer(config, task, self.get_temp_dir())
     if mixed_precision_dtype != 'float16':
-      self.assertIsInstance(trainer.optimizer, tf.keras.optimizers.SGD)
+      self.assertIsInstance(
+          trainer.optimizer,
+          (tf.keras.optimizers.SGD, tf.keras.optimizers.legacy.SGD))
     elif mixed_precision_dtype == 'float16' and loss_scale is None:
-      self.assertIsInstance(trainer.optimizer, tf.keras.optimizers.SGD)
+      self.assertIsInstance(
+          trainer.optimizer,
+          (tf.keras.optimizers.SGD, tf.keras.optimizers.legacy.SGD))
 
     metrics = trainer.train(tf.convert_to_tensor(5, dtype=tf.int32))
     self.assertIn('training_loss', metrics)