Fix Keras imports for optimizer algorithms

master · Feb 15, 2022 · 4814b0c · 4814b0c
1 parent 90c432a
commit 4814b0c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/tensorflow_riemopt/optimizers/constrained_rmsprop.py b/tensorflow_riemopt/optimizers/constrained_rmsprop.py
@@ -9,19 +9,19 @@
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import ops
 from tensorflow.python.keras import backend_config
-from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.keras.utils import generic_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.training import gen_training_ops
+from keras.optimizer_v2.optimizer_v2 import OptimizerV2
 
 from tensorflow_riemopt.variable import get_manifold
 
 
 @generic_utils.register_keras_serializable(name="ConstrainedRMSprop")
-class ConstrainedRMSprop(optimizer_v2.OptimizerV2):
+class ConstrainedRMSprop(OptimizerV2):
     """Optimizer that implements the RMSprop algorithm."""
 
     _HAS_AGGREGATE_GRAD = True

diff --git a/tensorflow_riemopt/optimizers/riemannian_adam.py b/tensorflow_riemopt/optimizers/riemannian_adam.py
@@ -6,19 +6,19 @@
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import ops
 from tensorflow.python.keras import backend_config
-from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.keras.utils import generic_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.training import gen_training_ops
+from keras.optimizer_v2.optimizer_v2 import OptimizerV2
 
 from tensorflow_riemopt.variable import get_manifold
 
 
 @generic_utils.register_keras_serializable(name="RiemannianAdam")
-class RiemannianAdam(optimizer_v2.OptimizerV2):
+class RiemannianAdam(OptimizerV2):
     """Optimizer that implements the Riemannian Adam algorithm."""
 
     _HAS_AGGREGATE_GRAD = True

diff --git a/tensorflow_riemopt/optimizers/riemannian_gradient_descent.py b/tensorflow_riemopt/optimizers/riemannian_gradient_descent.py
@@ -6,19 +6,19 @@
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import ops
 from tensorflow.python.keras import backend_config
-from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.keras.utils import generic_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.training import gen_training_ops
+from keras.optimizer_v2.optimizer_v2 import OptimizerV2
 
 from tensorflow_riemopt.variable import get_manifold
 
 
 @generic_utils.register_keras_serializable(name="RiemannianSGD")
-class RiemannianSGD(optimizer_v2.OptimizerV2):
+class RiemannianSGD(OptimizerV2):
     """Optimizer that implements the Riemannian SGD algorithm."""
 
     _HAS_AGGREGATE_GRAD = True