Use legacy keras optimizer to be compatible with an incoming Keras op…

…timizer migration. PiperOrigin-RevId: 467992914
google-deepmind · Aug 16, 2022 · c92ac8e · c92ac8e
1 parent d1cd371
commit c92ac8e
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/sonnet/src/optimizers/momentum_test.py b/sonnet/src/optimizers/momentum_test.py
@@ -201,7 +201,10 @@ def make_optimizer(self, **kwargs):
     if "use_nesterov" in kwargs:
       kwargs["nesterov"] = kwargs["use_nesterov"]
       del kwargs["use_nesterov"]
-    return optimizer_tests.WrappedTFOptimizer(tf.optimizers.SGD(**kwargs))
+    if hasattr(tf.keras.optimizers, "legacy"):
+      return optimizer_tests.WrappedTFOptimizer(
+          tf.keras.optimizers.legacy.SGD(**kwargs))
+    return optimizer_tests.WrappedTFOptimizer(tf.keras.optimizers.SGD(**kwargs))
 
 
 if __name__ == "__main__":

diff --git a/sonnet/src/optimizers/rmsprop_test.py b/sonnet/src/optimizers/rmsprop_test.py
@@ -201,7 +201,11 @@ def make_optimizer(self, **kwargs):
     if "learning_rate" not in kwargs:
       kwargs["learning_rate"] = 0.1
     kwargs["rho"] = kwargs.pop("decay", 0.9)
-    return optimizer_tests.WrappedTFOptimizer(tf.optimizers.RMSprop(**kwargs))
+    if hasattr(tf.keras.optimizers, "legacy"):
+      return optimizer_tests.WrappedTFOptimizer(
+          tf.keras.optimizers.legacy.RMSprop(**kwargs))
+    return optimizer_tests.WrappedTFOptimizer(
+        tf.keras.optimizers.RMSprop(**kwargs))
 
 
 if __name__ == "__main__":

diff --git a/sonnet/src/optimizers/sgd_test.py b/sonnet/src/optimizers/sgd_test.py
@@ -82,7 +82,10 @@ class ReferenceSGDTest(SGDTest):
   def make_optimizer(self, *args, **kwargs):
     if "learning_rate" not in kwargs:
       kwargs["learning_rate"] = 3.
-    return optimizer_tests.WrappedTFOptimizer(tf.optimizers.SGD(**kwargs))
+    if hasattr(tf.keras.optimizers, "legacy"):
+      return optimizer_tests.WrappedTFOptimizer(
+          tf.keras.optimizers.legacy.SGD(**kwargs))
+    return optimizer_tests.WrappedTFOptimizer(tf.keras.optimizers.SGD(**kwargs))
 
 
 if __name__ == "__main__":