Repository: horovod/horovod

Issue:
TF 2 issue introduced in commit 2be9d149
@alsrgv : Your last commit: 2be9d1492a1483a9d1f38343261d9350bdf1dcd5 introduces an issue with Tensorflow 2.

```python
File "/workspace/mask_rcnn/mask_rcnn_model.py", line 459, in _model_fn
    train_op = optimizer.apply_gradients(grads_and_vars, global_step=global_step)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/training/experimental/loss_scale_optimizer.py", line 183, in apply_gradients
    self._distributed_apply, args=(grads_and_vars, global_step, name))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1940, in merge_call
    return self._merge_call(merge_fn, args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1947, in _merge_call
    return merge_fn(self._strategy, *args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/training/experimental/loss_scale_optimizer.py", line 218, in _distributed_apply
    control_flow_ops.no_op)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/smart_cond.py", line 59, in smart_cond
    name=name)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/control_flow_ops.py", line 1174, in cond
    return cond_v2.cond_v2(pred, true_fn, false_fn, name)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/cond_v2.py", line 84, in cond_v2
    op_return_value=pred)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py", line 915, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/training/experimental/loss_scale_optimizer.py", line 215, in apply_fn
    name + '-wrapped')
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/training/experimental/loss_scale_optimizer.py", line 226, in _apply_gradients
    args=(grads_and_vars, global_step, name))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1810, in call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 2155, in _call_for_each_replica
    return fn(*args, **kwargs)
TypeError: apply_gradients() takes 2 positional arguments but 4 were given
```

If I rollback to commit 61e62f7a0a6a73a567c96e61bb6edbe861c5181a, everything is fine ;)


Discussion:

Resolving patch:
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -28,7 +28,7 @@ def __init__(self, name, device_dense, device_sparse, compression, sparse_as_den
                 name, device_dense, device_sparse, compression, sparse_as_dense)
             super(self.__class__, self).__init__(**config)
 
-        def apply_gradients(self, grads_and_vars, **kwargs):
+        def apply_gradients(self, grads_and_vars, *args, **kwargs):
             """Apply gradients to provided variables.
 
             See Optimizer.apply_gradients() for more info.
@@ -40,7 +40,7 @@ def apply_gradients(self, grads_and_vars, **kwargs):
                 grads, vars = zip(*grads_and_vars)
                 avg_grads = self._allreduce_grads(grads)
                 grads_and_vars = list(zip(avg_grads, vars))
-            return super(self.__class__, self).apply_gradients(grads_and_vars, **kwargs)
+            return super(self.__class__, self).apply_gradients(grads_and_vars, *args, **kwargs)
 
     class _DistributedOptimizerWithGetGradients(keras.optimizers.Optimizer):
         def __init__(self, name, device_dense, device_sparse, compression, sparse_as_dense,
diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py
--- a/horovod/tensorflow/__init__.py
+++ b/horovod/tensorflow/__init__.py
@@ -243,7 +243,7 @@ def __init__(self, optimizer, name=None, use_locking=False, device_dense='',
             self._allreduce_grads = _make_allreduce_grads_fn(
                 name, device_dense, device_sparse, compression, sparse_as_dense)
 
-        def apply_gradients(self, grads_and_vars, **kwargs):
+        def apply_gradients(self, grads_and_vars, *args, **kwargs):
             """Apply gradients to provided variables.
 
             See Optimizer.apply_gradients() for more info.
@@ -255,7 +255,7 @@ def apply_gradients(self, grads_and_vars, **kwargs):
                 grads, vars = zip(*grads_and_vars)
                 avg_grads = self._allreduce_grads(grads)
                 grads_and_vars = list(zip(avg_grads, vars))
-            return self._optimizer.apply_gradients(grads_and_vars, **kwargs)
+            return self._optimizer.apply_gradients(grads_and_vars, *args, **kwargs)
 
         def compute_gradients(self, *args, **kwargs):
             """Calls this same method on the underlying optimizer."""


Test patch: