Repository: tensorflow/addons

Issue:
Missing argument in apply_gradients() in AdamW optimizer
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
- TensorFlow version and how it was installed (source or binary): tf-nightly 2.2.0-dev20200309 (pip install)
- TensorFlow-Addons version and how it was installed (source or binary): 0.8.3 (pip install)
- Python version: 3.6
- Is GPU used? (yes/no): yes

**Describe the bug**

When running model.compile() with AdamW optimizer, a type error is thrown saying:
``apply_gradients() got an unexpected keyword argument 'all_reduce_sum_gradients'``

This can be fixed by adding in the argument to ``apply_gradients()`` in ``tensorflow_addons/optimizers/weight_decay_optimizers.py``

**Code to reproduce the issue**

https://colab.research.google.com/drive/1A6X8yYii5M8BDqwAFvoFglrTIqWvwQLm

**Other info / logs**

TypeError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:503 train_function  *
        outputs = self.distribute_strategy.experimental_run_v2(
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:920 experimental_run_v2  **
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2254 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2615 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:473 train_step  **
        _minimize(tape, self.optimizer, loss, self.trainable_variables)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:1737 _minimize
        all_reduce_sum_gradients=False)

    TypeError: apply_gradients() got an unexpected keyword argument 'all_reduce_sum_gradients'



Discussion:
You code works with tensorflow==2.1.0 and tfa-nightly.

We currently target the stable version of tensorflow (2.1.0) to, as you can see, keep the sanity of our devs. So we don't expect everything to work with tf-nightly.

But what you are reporting is concerning. Either :

* the bug comes from tensorflow and they're not being backward compatible there
* We use some undocumented/private API in AdamW and we should remove them.

In all cases let's keep this issue open.
@philjd could you look into it? It would be nice to fix it before the 0.9 release.
Thanks for flagging this. I just checked, the signature of the apply_gradients function has changed indeed. PR #1181 was an attempt to fix this but broke backward compatibility. I'll prepare a PR in a minute.
Thanks a lot. One thing I don't understand is that if when doing the subclassing, we used only public APIs, the breakage shouldn't happend right?

https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer#write_a_customized_optimizer_2

It seems we shouldn't override `apply_gradients`.
The problem is that we need to pass a list with the variables to decay. I put this into apply_gradients and not into e.g. constructor as the variables to decay are also passed there.

Atm I don't see a backward compatible way of getting rid of overwriting the apply_gradients function. 
I think it's easier to add **kwargs to the function signature, allowing us to be backward compatible.
Yep you're right. For now, let's fix the compatibility with TF 2.2.x with **kwargs. Then we can look into how to use only public API to reimplement it. If we don't do it at some point, we'll have to update this file at every TF release haha
Hm, technically we're just wrapping a public API, which seems perfectly valid to me.
I'd argue that the problem is on the Keras side, with an optimizer wrapping the public API in an incompatible way, requiring special treatment in the training loop [here](https://github.com/tensorflow/tensorflow/blob/v2.2.0-rc1/tensorflow/python/keras/engine/training.py#L1735).


We shouldn't override methods which are public and not meant to be overwritten. For example, if method x is public and not meant to be overwritten, and we overwrite it anyway, adding a new argument y, it's working for the moment. Then, since x is a public API, the TF team can also add a new argument z. Then we're in conflict on the API level with what does tensorflow.

But yeah the piece of code you're showing is unfortunate. There are some issues on both side. We can discuss at the meeting today if using composition over subclassing is possible, and if yes, then we may want to go this way and break backward compatibility.
We should discuss this one for sure. 
Hm, I just realized I took a look at the wrong snapshot for optimizer (2.1) and didn't get the full picture, sorry about that.
I guess we can get a long way with **kwargs (I think it's likely the TF team will only add new arguments with a default, which wouldn't require changes here) but I'd be happy to approach this differently if backward compatibility is less of a concern.

Resolving patch:
diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py
--- a/tensorflow_addons/optimizers/weight_decay_optimizers.py
+++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py
@@ -126,7 +126,7 @@ def minimize(self, loss, var_list, grad_loss=None, name=None, decay_var_list=Non
         )
         return super().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):
+    def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None, **kwargs):
         """Apply gradients to variables.
 
         This is the second part of `minimize()`. It returns an `Operation` that
@@ -138,6 +138,9 @@ def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None):
                 name passed to the `Optimizer` constructor.
             decay_var_list: Optional list of variables to be decayed. Defaults
                 to all variables in var_list.
+            **kwargs: Additional arguments to pass to the base optimizer's
+                apply_gradient method, e.g., TF2.2 added an argument
+                `all_reduce_sum_gradients`.
         Returns:
             An `Operation` that applies the specified gradients.
         Raises:
@@ -147,7 +150,7 @@ def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None):
         self._decay_var_list = (
             set([_ref(v) for v in decay_var_list]) if decay_var_list else False
         )
-        return super().apply_gradients(grads_and_vars, name=name)
+        return super().apply_gradients(grads_and_vars, name=name, **kwargs)
 
     def _decay_weights_op(self, var):
         if not self._decay_var_list or _ref(var) in self._decay_var_list:


Test patch:
diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
--- a/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
+++ b/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
@@ -273,6 +273,15 @@ def testBasicDecayVarList(self):
             weight_decay=WEIGHT_DECAY,
         )
 
+    def testKerasFit(self):
+        """Check if calling model.fit works."""
+        model = tf.keras.models.Sequential([tf.keras.layers.Dense(2)])
+        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+        optimizer = self.optimizer(learning_rate=1e-4, weight_decay=1e-4)
+        model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])
+        X, y = np.random.uniform(size=(2, 4, 1))
+        model.fit(X, y, epochs=1)
+
 
 @test_utils.run_all_in_graph_and_eager_modes
 class SGDWTest(OptimizerTestBase):