fix tests for tf 2.11 (#2783)

Author: chenmoneygithub

tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py

@@ -13,10 +13,16 @@
 EPOCHS = 5
 
 
+def get_legacy_sgd(learning_rate):
+    if hasattr(tf.keras.optimizers, "legacy"):
+        return tf.keras.optimizers.legacy.SGD(learning_rate)
+    return tf.keras.optimizers.SGD(learning_rate)
+
+
 def get_data_and_model(optimizer="moving_avg"):
     x = tf.random.normal([TRAIN_SAMPLES, INPUT_DIM])
     y = tf.random.normal([TRAIN_SAMPLES, NUM_CLASSES])
-    moving_avg = MovingAverage(tf.keras.optimizers.SGD(lr=2.0), average_decay=0.5)
+    moving_avg = MovingAverage(get_legacy_sgd(2.0), average_decay=0.5)
     if optimizer == "moving_avg":
         optimizer = moving_avg
     inputs = keras.layers.Input(INPUT_DIM)
@@ -199,7 +205,7 @@ def test_invalid_save_freq(tmp_path):
 
 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)
+    moving_avg = MovingAverage(get_legacy_sgd(2.0), average_decay=0.5)
     optimizer = tf.keras.mixed_precision.LossScaleOptimizer(moving_avg)
     x, y, model = get_data_and_model(optimizer)
     save_freq = "epoch"

tensorflow_addons/optimizers/tests/moving_average_test.py

@@ -222,9 +222,12 @@ def test_start_step():
     grads0 = tf.constant([0.1, 0.1])
     grads_and_vars = [(grads0, var0)]
 
-    opt = MovingAverage(
-        tf.keras.optimizers.SGD(lr=1.0), average_decay=0.5, start_step=1
-    )
+    if hasattr(tf.keras.optimizers, "legacy"):
+        sgd_opt = tf.keras.optimizers.legacy.SGD(lr=1.0)
+    else:
+        sgd_opt = tf.keras.optimizers.SGD(lr=1.0)
+
+    opt = MovingAverage(sgd_opt, average_decay=0.5, start_step=1)
 
     opt.apply_gradients(grads_and_vars)
 

tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py

@@ -431,9 +431,7 @@ def test_optimizer_basic(dtype, optimizer):
     "optimizer",
     [
         weight_decay_optimizers.SGDW,
-        weight_decay_optimizers.extend_with_decoupled_weight_decay(
-            tf.keras.optimizers.SGD
-        ),
+        weight_decay_optimizers.extend_with_decoupled_weight_decay(optimizer_class),
     ],
 )
 @pytest.mark.parametrize("dtype", [tf.half, tf.float32, tf.float64])

tensorflow_addons/optimizers/weight_decay_optimizers.py

@@ -371,8 +371,16 @@ def __init__(
     return OptimizerWithDecoupledWeightDecay
 
 
+if hasattr(tf.keras.optimizers, "legacy"):
+    ADAM_CLASS = tf.keras.optimizers.legacy.Adam
+    SGD_CLASS = tf.keras.optimizers.legacy.SGD
+else:
+    ADAM_CLASS = tf.keras.optimizers.Adam
+    SGD_CLASS = tf.keras.optimizers.SGD
+
+
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class SGDW(DecoupledWeightDecayExtension, tf.keras.optimizers.SGD):
+class SGDW(DecoupledWeightDecayExtension, SGD_CLASS):
     """Optimizer that implements the Momentum algorithm with weight_decay.
 
     This is an implementation of the SGDW optimizer described in "Decoupled
@@ -450,12 +458,6 @@ def __init__(
         )
 
 
-if hasattr(tf.keras.optimizers, "legacy"):
-    ADAM_CLASS = tf.keras.optimizers.legacy.Adam
-else:
-    ADAM_CLASS = tf.keras.optimizers.Adam
-
-
 @tf.keras.utils.register_keras_serializable(package="Addons")
 class AdamW(DecoupledWeightDecayExtension, ADAM_CLASS):
     """Optimizer that implements the Adam algorithm with weight decay.