Point optimizer to tf.keras.optimizer.legacy.Optimizer to be compatib… (#2706)

* Point optimizer to tf.keras.optimizer.legacy.Optimizer to be compatible with Keras optimizer migration

* small fix

* add version control

* small fix

* Update discriminative_layer_training.py

* fix version control

* small fix

* move optimizer class to __init__.py

* small fix

* fix problems

* small fix

* Rename BaseOptimizer to KerasLegacyOptimizer

* exclude keras optimizer from type check

* fix import

Author: chenmoneygithub

tensorflow_addons/optimizers/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Additional optimizers that conform to Keras API."""
 
+from tensorflow_addons.optimizers.constants import KerasLegacyOptimizer
 from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper
 from tensorflow_addons.optimizers.conditional_gradient import ConditionalGradient
 from tensorflow_addons.optimizers.cyclical_learning_rate import CyclicalLearningRate

tensorflow_addons/optimizers/adabelief.py

@@ -17,11 +17,12 @@
 import tensorflow as tf
 from tensorflow_addons.utils.types import FloatTensorLike
 
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
 from typing import Union, Callable, Dict
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class AdaBelief(tf.keras.optimizers.Optimizer):
+class AdaBelief(KerasLegacyOptimizer):
     """Variant of the Adam optimizer.
 
     It achieves fast convergence as Adam and generalization comparable to SGD.

tensorflow_addons/optimizers/average_wrapper.py

@@ -17,12 +17,12 @@
 import warnings
 
 import tensorflow as tf
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
 from tensorflow_addons.utils import types
-
 from typeguard import typechecked
 
 
-class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):
+class AveragedOptimizerWrapper(KerasLegacyOptimizer, metaclass=abc.ABCMeta):
     @typechecked
     def __init__(
         self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs
@@ -32,9 +32,12 @@ def __init__(
         if isinstance(optimizer, str):
             optimizer = tf.keras.optimizers.get(optimizer)
 
-        if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        if not isinstance(
+            optimizer, (tf.keras.optimizers.Optimizer, KerasLegacyOptimizer)
+        ):
             raise TypeError(
-                "optimizer is not an object of tf.keras.optimizers.Optimizer"
+                "optimizer is not an object of tf.keras.optimizers.Optimizer "
+                "or tf.keras.optimizers.legacy.Optimizer (if you have tf version >= 2.9.0)."
             )
 
         self._optimizer = optimizer

tensorflow_addons/optimizers/cocob.py

@@ -17,9 +17,11 @@
 from typeguard import typechecked
 import tensorflow as tf
 
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
+
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class COCOB(tf.keras.optimizers.Optimizer):
+class COCOB(KerasLegacyOptimizer):
     """Optimizer that implements COCOB Backprop Algorithm
 
         Reference:

tensorflow_addons/optimizers/conditional_gradient.py

@@ -15,14 +15,15 @@
 """Conditional Gradient optimizer."""
 
 import tensorflow as tf
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
 from tensorflow_addons.utils.types import FloatTensorLike
 
 from typeguard import typechecked
 from typing import Union, Callable
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class ConditionalGradient(tf.keras.optimizers.Optimizer):
+class ConditionalGradient(KerasLegacyOptimizer):
     """Optimizer that implements the Conditional Gradient optimization.
 
     This optimizer helps handle constraints well.

tensorflow_addons/optimizers/constants.py

@@ -0,0 +1,21 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import importlib
+import tensorflow as tf
+
+if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None:
+    KerasLegacyOptimizer = tf.keras.optimizers.legacy.Optimizer
+else:
+    KerasLegacyOptimizer = tf.keras.optimizers.Optimizer

tensorflow_addons/optimizers/cyclical_learning_rate.py

@@ -58,7 +58,7 @@ def __init__(
         ```
 
         You can pass this schedule directly into a
-        `tf.keras.optimizers.Optimizer` as the learning rate.
+        `tf.keras.optimizers.legacy.Optimizer` as the learning rate.
 
         Args:
             initial_learning_rate: A scalar `float32` or `float64` `Tensor` or
@@ -146,7 +146,7 @@ def __init__(
         ```
 
         You can pass this schedule directly into a
-        `tf.keras.optimizers.Optimizer` as the learning rate.
+        `tf.keras.optimizers.legacy.Optimizer` as the learning rate.
 
         Args:
             initial_learning_rate: A scalar `float32` or `float64` `Tensor` or
@@ -215,7 +215,7 @@ def __init__(
         ```
 
         You can pass this schedule directly into a
-        `tf.keras.optimizers.Optimizer` as the learning rate.
+        `tf.keras.optimizers.legacy.Optimizer` as the learning rate.
 
         Args:
             initial_learning_rate: A scalar `float32` or `float64` `Tensor` or
@@ -286,7 +286,7 @@ def __init__(
         ```
 
         You can pass this schedule directly into a
-        `tf.keras.optimizers.Optimizer` as the learning rate.
+        `tf.keras.optimizers.legacy.Optimizer` as the learning rate.
 
         Args:
             initial_learning_rate: A scalar `float32` or `float64` `Tensor` or

tensorflow_addons/optimizers/discriminative_layer_training.py

@@ -17,14 +17,16 @@
 from typing import List, Union
 
 import tensorflow as tf
+
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
 from typeguard import typechecked
 
 from keras import backend
 from keras.utils import tf_utils
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class MultiOptimizer(tf.keras.optimizers.Optimizer):
+class MultiOptimizer(KerasLegacyOptimizer):
     """Multi Optimizer Wrapper for Discriminative Layer Training.
 
     Creates a wrapper around a set of instantiated optimizer layer pairs.
@@ -33,7 +35,7 @@ class MultiOptimizer(tf.keras.optimizers.Optimizer):
     Each optimizer will optimize only the weights associated with its paired layer.
     This can be used to implement discriminative layer training by assigning
     different learning rates to each optimizer layer pair.
-    `(tf.keras.optimizers.Optimizer, List[tf.keras.layers.Layer])` pairs are also supported.
+    `(tf.keras.optimizers.legacy.Optimizer, List[tf.keras.layers.Layer])` pairs are also supported.
     Please note that the layers must be instantiated before instantiating the optimizer.
 
     Args:
@@ -149,7 +151,7 @@ def get_config(self):
     @classmethod
     def create_optimizer_spec(
         cls,
-        optimizer: tf.keras.optimizers.Optimizer,
+        optimizer: KerasLegacyOptimizer,
         layers_or_model: Union[
             tf.keras.Model,
             tf.keras.Sequential,

tensorflow_addons/optimizers/lamb.py

@@ -24,12 +24,13 @@
 from typeguard import typechecked
 
 import tensorflow as tf
+from tensorflow_addons.optimizers import KerasLegacyOptimizer
 from tensorflow_addons.utils.types import FloatTensorLike
 from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class LAMB(tf.keras.optimizers.Optimizer):
+class LAMB(KerasLegacyOptimizer):
     """Optimizer that implements the Layer-wise Adaptive Moments (LAMB).
 
     See paper [Large Batch Optimization for Deep Learning: Training BERT

tensorflow_addons/optimizers/lazy_adam.py

@@ -20,15 +20,22 @@
 original Adam algorithm, and may lead to different empirical results.
 """
 
+import importlib
 import tensorflow as tf
 from tensorflow_addons.utils.types import FloatTensorLike
 
 from typeguard import typechecked
 from typing import Union, Callable
 
 
+if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None:
+    adam_optimizer_class = tf.keras.optimizers.legacy.Adam
+else:
+    adam_optimizer_class = tf.keras.optimizers.Adam
+
+
 @tf.keras.utils.register_keras_serializable(package="Addons")
-class LazyAdam(tf.keras.optimizers.Adam):
+class LazyAdam(adam_optimizer_class):
     """Variant of the Adam optimizer that handles sparse updates more
     efficiently.