Delete optimizer_test_base.py Remove keras object registration in the

factory function.

Author: PhilJd

tensorflow_addons/optimizers/README.md

@@ -27,7 +27,6 @@ must:
    `@run_in_graph_and_eager_modes` (for test method)
    or `run_all_in_graph_and_eager_modes` (for TestCase subclass)
    decorator.
- * Consider inheriting from `OptimizerTestBase`.
  * Add a `py_test` to this sub-package's BUILD file.
 
 #### Documentation Requirements

tensorflow_addons/optimizers/optimizer_test_base.py

@@ -125,8 +125,11 @@ def minimize(self,
             ValueError: If some of the variables are not `Variable` objects.
         """
         self._decay_var_list = set(decay_var_list) if decay_var_list else False
-        return super(DecoupledWeightDecayExtension, self).minimize(
-            loss, var_list=var_list, grad_loss=grad_loss, name=name)
+        return super(DecoupledWeightDecayExtension,
+                     self).minimize(loss,
+                                    var_list=var_list,
+                                    grad_loss=grad_loss,
+                                    name=name)
 
     def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None):
         """Apply gradients to variables.
@@ -149,8 +152,8 @@ def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None):
             ValueError: If none of the variables have gradients.
         """
         self._decay_var_list = set(decay_var_list) if decay_var_list else False
-        return super(DecoupledWeightDecayExtension, self).apply_gradients(
-            grads_and_vars, name=name)
+        return super(DecoupledWeightDecayExtension,
+                     self).apply_gradients(grads_and_vars, name=name)
 
     def _decay_weights_op(self, var):
         if not self._decay_var_list or var in self._decay_var_list:
@@ -161,8 +164,8 @@ def _decay_weights_op(self, var):
 
     def _decay_weights_sparse_op(self, var, indices):
         if not self._decay_var_list or var in self._decay_var_list:
-            update = (-self._get_hyper('weight_decay', var.dtype) * tf.gather(
-                var, indices))
+            update = (-self._get_hyper('weight_decay', var.dtype) *
+                      tf.gather(var, indices))
             return self._resource_scatter_add(var, indices, update)
         return tf.no_op()
 
@@ -226,17 +229,19 @@ def extend_with_decoupled_weight_decay(base_optimizer):
 
     optimizer = tfa.optimizers.AdamW(learning_rate=lr, weight_decay=wd)
     ```
+    
+    Note: you might want to register your own custom optimizer using
+    `tf.keras.utils.get_custom_objects()`.
 
     Args:
-      base_optimizer: An optimizer class that inherits from
-          tf.optimizers.Optimizer.
+        base_optimizer: An optimizer class that inherits from
+            tf.optimizers.Optimizer.
 
     Returns:
-      A new optimizer class that inherits from DecoupledWeightDecayExtension
-      and base_optimizer.
+        A new optimizer class that inherits from DecoupledWeightDecayExtension
+        and base_optimizer.
     """
 
-    @keras_utils.register_keras_custom_object
     class OptimizerWithDecoupledWeightDecay(DecoupledWeightDecayExtension,
                                             base_optimizer):
         """Base_optimizer with decoupled weight decay.
@@ -255,8 +260,8 @@ class OptimizerWithDecoupledWeightDecay(DecoupledWeightDecayExtension,
 
         def __init__(self, weight_decay, *args, **kwargs):
             # super delegation is necessary here
-            super(OptimizerWithDecoupledWeightDecay, self).__init__(
-                weight_decay, *args, **kwargs)
+            super(OptimizerWithDecoupledWeightDecay,
+                  self).__init__(weight_decay, *args, **kwargs)
 
     return OptimizerWithDecoupledWeightDecay
 
@@ -326,13 +331,12 @@ def __init__(self,
                 of learning rate. `lr` is included for backward compatibility,
                 recommended to use `learning_rate` instead.
         """
-        super(SGDW, self).__init__(
-            weight_decay,
-            learning_rate=learning_rate,
-            momentum=momentum,
-            nesterov=nesterov,
-            name=name,
-            **kwargs)
+        super(SGDW, self).__init__(weight_decay,
+                                   learning_rate=learning_rate,
+                                   momentum=momentum,
+                                   nesterov=nesterov,
+                                   name=name,
+                                   **kwargs)
 
 
 @keras_utils.register_keras_custom_object
@@ -412,12 +416,11 @@ def __init__(self,
                 of learning rate. `lr` is included for backward compatibility,
                 recommended to use `learning_rate` instead.
         """
-        super(AdamW, self).__init__(
-            weight_decay,
-            learning_rate=learning_rate,
-            beta_1=beta_1,
-            beta_2=beta_2,
-            epsilon=epsilon,
-            amsgrad=amsgrad,
-            name=name,
-            **kwargs)
+        super(AdamW, self).__init__(weight_decay,
+                                    learning_rate=learning_rate,
+                                    beta_1=beta_1,
+                                    beta_2=beta_2,
+                                    epsilon=epsilon,
+                                    amsgrad=amsgrad,
+                                    name=name,
+                                    **kwargs)