We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Bug in global_step update when local gradient aggregation is used
**Environment:**
1. Framework: TensorFlow
2. Framework version: 1.15.3
3. Horovod version: 0.21.3
4. MPI version: 3.1.6
5. CUDA version: 11.0.194
6. NCCL version: 2.7.6
7. Python version: 3.6.9
8. Spark / PySpark version:
9. OS and version: 
10. GCC version:
11. CMake version:

**Bug report:**
There seems to be a bug in how the `global_step` variable is updated when using local gradient aggregation with `backward_passes_per_step>1`. 

Here is a simple example to reproduce the issue:
```
import tensorflow as tf
import horovod.tensorflow as hvd

tf.disable_eager_execution()
LR=0.01
# data/label is always just 1
data = tf.constant([[1]], dtype=tf.float32)
# one-variable linear model, no bias (y = w*x)
#
# gradient update is hand-calculatable if needed for testing:
#
# model:            yhat = w*x
# mse loss:         L = (ytrue - w*x)**2
# gradient:         dL/dw = -2*x*(ytrue - `w*x)`
# let ytrue = x:    dL/dw = -2*x*x*(1-w)
def predict(x):
    dense = tf.layers.Dense(1, use_bias=False, kernel_initializer=tf.zeros_initializer())
    return dense(x)
pred = predict(data)
loss = tf.losses.mean_squared_error(data, pred)
step = tf.contrib.framework.get_or_create_global_step()
optimizer = tf.train.GradientDescentOptimizer(LR)
hvd.init()
optimizer = hvd.DistributedOptimizer(optimizer, backward_passes_per_step=2, average_aggregated_gradients=True)

with tf.control_dependencies([tf.print("step:", step, "loss:", loss)]):
    train_op = optimizer.minimize(loss, global_step=step)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    # training loop
    for i in range(10):
        sess.run([train_op])
```

Which gives this output:
```
step: 0 loss: 1
step: 0 loss: 1
step: 0 loss: 1
step: 2 loss: 0.960400045
step: 2 loss: 0.960400045
step: 4 loss: 0.922368109
step: 4 loss: 0.922368109
step: 6 loss: 0.885842443
step: 6 loss: 0.885842443
step: 8 loss: 0.850763
step: 8 loss: 0.850763
```




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `horovod/tensorflow/gradient_aggregation.py`
Content:
```
1 import tensorflow as tf
2 
3 
4 def apply_op_to_not_none_tensors(tensor_op, tensors, *args):
5     return [
6         tensor_op(
7             tensor,
8             *args
9         ) if tensor is not None else tensor for tensor in tensors]
10 
11 
12 def get_not_none_from_list(tensor_list):
13     return [x for x in tensor_list if x is not None]
14 
15 
16 class LocalGradientAggregationHelper:
17     """
18     LocalGradientAggregationHelper aggregates gradient updates locally,
19     and communicates the updates across machines only once every
20     backward_passes_per_step. Only supports graph mode execution.
21     """
22 
23     _OPTIMIZER_TYPE_KERAS = "optimizer_type_keras"
24     _OPTIMIZER_TYPE_LEGACY = "optimizer_type_legacy"
25 
26     def __init__(
27             self,
28             backward_passes_per_step,
29             allreduce_func,
30             sparse_as_dense,
31             average_aggregated_gradients,
32             rank,
33             optimizer_type):
34         self._allreduce_grads = allreduce_func
35 
36         # backward_passes_per_step controls how often gradient updates are
37         # synchronized.
38         self.backward_passes_per_step = backward_passes_per_step
39         if self.backward_passes_per_step <= 0:
40             raise ValueError("backward_passes_per_step must be > 0")
41 
42         # average_aggregated_gradients controls whether gradient updates that are
43         # aggregated, should be divided by `backward_passes_per_step`.
44         self.average_aggregated_gradients = average_aggregated_gradients
45 
46         # This is going to be [N] data structure holding the aggregated gradient updates
47         # N is the number of parameters.
48         self.locally_aggregated_grads = []
49 
50         # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`
51         # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is
52         # equal to 0.
53         self.counter = None
54 
55         self.sparse_as_dense = sparse_as_dense
56         self.rank = rank
57         self.optimizer_type = optimizer_type
58 
59         # Contains the mapping of indexes of grad updates that are not None to their index in
60         # locally_aggregated_grads which only contains not None gradients. When performing
61         # gradient aggregation we have to remove them from the list of grads prior to passing
62         # the list into a tf.cond().
63         self.not_none_indexes = {}
64         self.num_none_grad_updates = 0
65 
66     def _init_aggregation_vars(self, grads):
67         """
68         Initializes the counter that is used when to communicate and aggregate gradients
69         and the tensorflow variables that store the locally aggregated gradients.
70         """
71         variable_scope_name = "aggregation_variables_" + str(self.rank)
72         with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):
73             self.counter = tf.compat.v1.get_variable(
74                 "aggregation_counter", shape=(), dtype=tf.int32,
75                 trainable=False, initializer=tf.compat.v1.zeros_initializer(),
76                 collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES],
77             )
78             for idx, grad in enumerate(grads):
79                 # Handle IndexedSlices.
80                 if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
81                     grad = tf.convert_to_tensor(grad)
82                 elif isinstance(grad, tf.IndexedSlices):
83                     raise ValueError(
84                         "IndexedSlices are not supported when "
85                         "`backward_passes_per_step` > 1 and "
86                         "`sparse_as_dense` is False."
87                     )
88 
89                 # Handle grads that are None.
90                 if grad is None:
91                     self.num_none_grad_updates += 1
92                     continue
93                 self.not_none_indexes[idx] = len(self.locally_aggregated_grads)
94 
95                 # Create shadow variable.
96                 grad_aggregation_variable_name = str(idx)
97                 zero_grad = tf.zeros(shape=grad.get_shape().as_list(), dtype=grad.dtype)
98                 grad_aggregation_variable = tf.compat.v1.get_variable(
99                     grad_aggregation_variable_name,
100                     trainable=False,
101                     initializer=zero_grad,
102                     collections=[
103                         tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
104                         "aggregating_collection"],
105                 )
106                 self.locally_aggregated_grads.append(grad_aggregation_variable)
107             assert len(self.locally_aggregated_grads) + \
108                 self.num_none_grad_updates == len(grads)
109 
110         # We expect to get a `sess` when we need to manually do a `sess.run(...)`
111         # for the variables to be initialized. This is the `tf.keras`
112         # optimizers.
113         if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:
114             session = tf.compat.v1.keras.backend.get_session(op_input_list=())
115             vars_init_op = tf.compat.v1.variables_initializer(
116                 [self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]
117             )
118             session.run(vars_init_op)
119 
120     def _clear_grads(self):
121         clear_ops_list = []
122         for idx, grad_aggregator in enumerate(self.locally_aggregated_grads):
123             clear_op = grad_aggregator.assign(grad_aggregator.initial_value)
124             clear_ops_list.append(clear_op)
125         return tf.group(*clear_ops_list)
126 
127     def _aggregate_grads(self, grads):
128         aggregation_ops_list = []
129         grads = get_not_none_from_list(grads)
130         assert len(grads) == len(self.locally_aggregated_grads)
131 
132         # Apply new gradient updates to the local copy.
133         for idx, grad in enumerate(grads):
134             if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
135                 grad = tf.convert_to_tensor(grad)
136 
137             updated_grad_aggregator = self.locally_aggregated_grads[idx].assign_add(
138                 grad)
139             aggregation_ops_list.append(updated_grad_aggregator)
140 
141         return aggregation_ops_list
142 
143     def _allreduce_grads_helper(self, grads, vars):
144         # Read in latest variables values.
145         aggregated_grads = []
146         aggregation_read_ops_list = []
147         for idx, locally_aggregated_grad in enumerate(
148                 self.locally_aggregated_grads):
149             aggregated_grads.append(locally_aggregated_grad.read_value())
150             aggregation_read_ops_list.append(aggregated_grads[idx])
151         aggregation_read_ops = tf.group(*aggregation_read_ops_list)
152 
153         with tf.control_dependencies([aggregation_read_ops]):
154             averaged_gradients = self._allreduce_grads(aggregated_grads, vars)
155 
156             # Reset counter.
157             with tf.control_dependencies([g.op for g in averaged_gradients if g is not None]):
158                 reset_op = self.counter.assign(
159                     tf.constant(0), use_locking=True)
160 
161             # Divide by backward_passes_per_step if
162             # average_aggregated_gradients is True.
163             with tf.control_dependencies([reset_op]):
164                 gradient_divisor = self.backward_passes_per_step if \
165                     self.average_aggregated_gradients else 1
166 
167                 averaged_gradients = apply_op_to_not_none_tensors(
168                     tf.divide,
169                     averaged_gradients,
170                     gradient_divisor,
171                 )
172                 return averaged_gradients
173 
174     def compute_gradients(self, grads, vars):
175         """
176         Applies the new gradient updates the locally aggregated gradients, and
177         performs cross-machine communication every backward_passes_per_step
178         times it is called.
179         """
180         self._init_aggregation_vars(grads)
181 
182         # Clear the locally aggregated gradients when the counter is at zero.
183         clear_op = tf.cond(
184             pred=tf.equal(self.counter, 0),
185             true_fn=lambda: self._clear_grads(),
186             false_fn=tf.no_op
187         )
188 
189         # Add new gradients to the locally aggregated gradients.
190         with tf.control_dependencies([clear_op]):
191             aggregation_ops_list = self._aggregate_grads(grads)
192 
193         # Increment the counter once new gradients have been applied.
194         aggregation_ops = tf.group(*aggregation_ops_list)
195         with tf.control_dependencies([aggregation_ops]):
196             update_counter = self.counter.assign_add(tf.constant(1))
197 
198         with tf.control_dependencies([update_counter]):
199             grads = get_not_none_from_list(grads)
200             assert len(grads) == len(self.locally_aggregated_grads)
201 
202             # Allreduce locally aggregated gradients when the counter is equivalent to
203             # `backward_passes_per_step`. This the condition is true, it also resets
204             # the counter back to 0.
205             allreduced_grads = tf.cond(
206                 tf.equal(self.counter, self.backward_passes_per_step),
207                 lambda: self._allreduce_grads_helper(grads, vars),
208                 lambda: grads,
209             )
210 
211             # Handle case where there is only one variable.
212             if not isinstance(allreduced_grads, (list, tuple)):
213                 allreduced_grads = (allreduced_grads,)
214             assert len(allreduced_grads) == len(self.locally_aggregated_grads)
215 
216             # Insert gradients that are None back in.
217             allreduced_grads = [
218                 allreduced_grads[self.not_none_indexes[idx]] if idx in self.not_none_indexes else None
219                 for idx in range(len(self.locally_aggregated_grads) + self.num_none_grad_updates)
220             ]
221             assert len(allreduced_grads) == len(
222                 self.locally_aggregated_grads) + self.num_none_grad_updates
223 
224         # If gradients have not been allreduced this batch, we return the gradients
225         # that were submitted as the updates (the input).
226         return allreduced_grads
227 
228     def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
229         """
230         Apply updates every backward_passes_per_step, which lines up with
231         the batches on which we communicated the locally aggregated gradients.
232         """
233         flattended_args0 = [item for tup in args[0] for item in tup]
234 
235         # Since we skip applying updates when the counter is not at zero we
236         # still want to increment the global step if it is being tracked
237         # (e.g., Tensorflow Estimators).
238         def increment_global_step_counter():
239             global_step_counter = tf.compat.v1.train.get_global_step()
240             if global_step_counter is None:
241                 return tf.no_op()
242             return global_step_counter.assign_add(
243                 tf.constant(1, dtype=tf.int64),
244                 use_locking=True,
245                 read_value=False
246             )
247 
248         # Increment global step on iterations where we don't call `apply_gradients()`.
249         cond_increment_global_step_counter = tf.cond(
250             pred=tf.equal(self.counter, 0),
251             true_fn=tf.no_op,
252             false_fn=increment_global_step_counter,
253         )
254         flattended_args0.append(cond_increment_global_step_counter)
255 
256         # If optimizer tracks iterations, we increment it on steps where we
257         # are not going to call `apply_gradients()`.
258         def increment_optimizer_iteration():
259             if hasattr(optimizer, "_iterations") and optimizer._iterations is not None:
260                 return optimizer._iterations.assign_add(1).op
261             return tf.no_op()
262 
263         with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):
264             return tf.cond(
265                 pred=tf.equal(self.counter, 0),
266                 true_fn=apply_grads_closure,
267                 false_fn=increment_optimizer_iteration,
268             )
269
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/horovod/tensorflow/gradient_aggregation.py b/horovod/tensorflow/gradient_aggregation.py
--- a/horovod/tensorflow/gradient_aggregation.py
+++ b/horovod/tensorflow/gradient_aggregation.py
@@ -232,6 +232,20 @@
         """
         flattended_args0 = [item for tup in args[0] for item in tup]
 
+        # If optimizer tracks iterations, we increment it on steps where we
+        # are not going to call `apply_gradients()`.
+        def increment_optimizer_iteration():
+            if hasattr(optimizer, "_iterations") and optimizer._iterations is not None:
+                return optimizer._iterations.assign_add(1).op
+            return tf.no_op()
+
+        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):
+            train_op = tf.cond(
+                pred=tf.equal(self.counter, 0),
+                true_fn=apply_grads_closure,
+                false_fn=increment_optimizer_iteration,
+            )
+
         # Since we skip applying updates when the counter is not at zero we
         # still want to increment the global step if it is being tracked
         # (e.g., Tensorflow Estimators).
@@ -245,24 +259,10 @@
                 read_value=False
             )
 
-        # Increment global step on iterations where we don't call `apply_gradients()`.
-        cond_increment_global_step_counter = tf.cond(
-            pred=tf.equal(self.counter, 0),
-            true_fn=tf.no_op,
-            false_fn=increment_global_step_counter,
-        )
-        flattended_args0.append(cond_increment_global_step_counter)
-
-        # If optimizer tracks iterations, we increment it on steps where we
-        # are not going to call `apply_gradients()`.
-        def increment_optimizer_iteration():
-            if hasattr(optimizer, "_iterations") and optimizer._iterations is not None:
-                return optimizer._iterations.assign_add(1).op
-            return tf.no_op()
-
-        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):
+        with tf.control_dependencies([train_op]):
+            # Increment global step on iterations where we don't call `apply_gradients()`.
             return tf.cond(
                 pred=tf.equal(self.counter, 0),
-                true_fn=apply_grads_closure,
-                false_fn=increment_optimizer_iteration,
+                true_fn=tf.no_op,
+                false_fn=increment_global_step_counter,
             )

{"golden_diff": "diff --git a/horovod/tensorflow/gradient_aggregation.py b/horovod/tensorflow/gradient_aggregation.py\n--- a/horovod/tensorflow/gradient_aggregation.py\n+++ b/horovod/tensorflow/gradient_aggregation.py\n@@ -232,6 +232,20 @@\n         \"\"\"\n         flattended_args0 = [item for tup in args[0] for item in tup]\n \n+        # If optimizer tracks iterations, we increment it on steps where we\n+        # are not going to call `apply_gradients()`.\n+        def increment_optimizer_iteration():\n+            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n+                return optimizer._iterations.assign_add(1).op\n+            return tf.no_op()\n+\n+        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):\n+            train_op = tf.cond(\n+                pred=tf.equal(self.counter, 0),\n+                true_fn=apply_grads_closure,\n+                false_fn=increment_optimizer_iteration,\n+            )\n+\n         # Since we skip applying updates when the counter is not at zero we\n         # still want to increment the global step if it is being tracked\n         # (e.g., Tensorflow Estimators).\n@@ -245,24 +259,10 @@\n                 read_value=False\n             )\n \n-        # Increment global step on iterations where we don't call `apply_gradients()`.\n-        cond_increment_global_step_counter = tf.cond(\n-            pred=tf.equal(self.counter, 0),\n-            true_fn=tf.no_op,\n-            false_fn=increment_global_step_counter,\n-        )\n-        flattended_args0.append(cond_increment_global_step_counter)\n-\n-        # If optimizer tracks iterations, we increment it on steps where we\n-        # are not going to call `apply_gradients()`.\n-        def increment_optimizer_iteration():\n-            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n-                return optimizer._iterations.assign_add(1).op\n-            return tf.no_op()\n-\n-        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):\n+        with tf.control_dependencies([train_op]):\n+            # Increment global step on iterations where we don't call `apply_gradients()`.\n             return tf.cond(\n                 pred=tf.equal(self.counter, 0),\n-                true_fn=apply_grads_closure,\n-                false_fn=increment_optimizer_iteration,\n+                true_fn=tf.no_op,\n+                false_fn=increment_global_step_counter,\n             )\n", "issue": "Bug in global_step update when local gradient aggregation is used\n**Environment:**\r\n1. Framework: TensorFlow\r\n2. Framework version: 1.15.3\r\n3. Horovod version: 0.21.3\r\n4. MPI version: 3.1.6\r\n5. CUDA version: 11.0.194\r\n6. NCCL version: 2.7.6\r\n7. Python version: 3.6.9\r\n8. Spark / PySpark version:\r\n9. OS and version: \r\n10. GCC version:\r\n11. CMake version:\r\n\r\n**Bug report:**\r\nThere seems to be a bug in how the `global_step` variable is updated when using local gradient aggregation with `backward_passes_per_step>1`. \r\n\r\nHere is a simple example to reproduce the issue:\r\n```\r\nimport tensorflow as tf\r\nimport horovod.tensorflow as hvd\r\n\r\ntf.disable_eager_execution()\r\nLR=0.01\r\n# data/label is always just 1\r\ndata = tf.constant([[1]], dtype=tf.float32)\r\n# one-variable linear model, no bias (y = w*x)\r\n#\r\n# gradient update is hand-calculatable if needed for testing:\r\n#\r\n# model:            yhat = w*x\r\n# mse loss:         L = (ytrue - w*x)**2\r\n# gradient:         dL/dw = -2*x*(ytrue - `w*x)`\r\n# let ytrue = x:    dL/dw = -2*x*x*(1-w)\r\ndef predict(x):\r\n    dense = tf.layers.Dense(1, use_bias=False, kernel_initializer=tf.zeros_initializer())\r\n    return dense(x)\r\npred = predict(data)\r\nloss = tf.losses.mean_squared_error(data, pred)\r\nstep = tf.contrib.framework.get_or_create_global_step()\r\noptimizer = tf.train.GradientDescentOptimizer(LR)\r\nhvd.init()\r\noptimizer = hvd.DistributedOptimizer(optimizer, backward_passes_per_step=2, average_aggregated_gradients=True)\r\n\r\nwith tf.control_dependencies([tf.print(\"step:\", step, \"loss:\", loss)]):\r\n    train_op = optimizer.minimize(loss, global_step=step)\r\nwith tf.Session() as sess:\r\n    sess.run(tf.global_variables_initializer())\r\n    sess.run(tf.local_variables_initializer())\r\n    # training loop\r\n    for i in range(10):\r\n        sess.run([train_op])\r\n```\r\n\r\nWhich gives this output:\r\n```\r\nstep: 0 loss: 1\r\nstep: 0 loss: 1\r\nstep: 0 loss: 1\r\nstep: 2 loss: 0.960400045\r\nstep: 2 loss: 0.960400045\r\nstep: 4 loss: 0.922368109\r\nstep: 4 loss: 0.922368109\r\nstep: 6 loss: 0.885842443\r\nstep: 6 loss: 0.885842443\r\nstep: 8 loss: 0.850763\r\nstep: 8 loss: 0.850763\r\n```\r\n\r\n\r\n\n", "before_files": [{"content": "import tensorflow as tf\n\n\ndef apply_op_to_not_none_tensors(tensor_op, tensors, *args):\n    return [\n        tensor_op(\n            tensor,\n            *args\n        ) if tensor is not None else tensor for tensor in tensors]\n\n\ndef get_not_none_from_list(tensor_list):\n    return [x for x in tensor_list if x is not None]\n\n\nclass LocalGradientAggregationHelper:\n    \"\"\"\n    LocalGradientAggregationHelper aggregates gradient updates locally,\n    and communicates the updates across machines only once every\n    backward_passes_per_step. Only supports graph mode execution.\n    \"\"\"\n\n    _OPTIMIZER_TYPE_KERAS = \"optimizer_type_keras\"\n    _OPTIMIZER_TYPE_LEGACY = \"optimizer_type_legacy\"\n\n    def __init__(\n            self,\n            backward_passes_per_step,\n            allreduce_func,\n            sparse_as_dense,\n            average_aggregated_gradients,\n            rank,\n            optimizer_type):\n        self._allreduce_grads = allreduce_func\n\n        # backward_passes_per_step controls how often gradient updates are\n        # synchronized.\n        self.backward_passes_per_step = backward_passes_per_step\n        if self.backward_passes_per_step <= 0:\n            raise ValueError(\"backward_passes_per_step must be > 0\")\n\n        # average_aggregated_gradients controls whether gradient updates that are\n        # aggregated, should be divided by `backward_passes_per_step`.\n        self.average_aggregated_gradients = average_aggregated_gradients\n\n        # This is going to be [N] data structure holding the aggregated gradient updates\n        # N is the number of parameters.\n        self.locally_aggregated_grads = []\n\n        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`\n        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is\n        # equal to 0.\n        self.counter = None\n\n        self.sparse_as_dense = sparse_as_dense\n        self.rank = rank\n        self.optimizer_type = optimizer_type\n\n        # Contains the mapping of indexes of grad updates that are not None to their index in\n        # locally_aggregated_grads which only contains not None gradients. When performing\n        # gradient aggregation we have to remove them from the list of grads prior to passing\n        # the list into a tf.cond().\n        self.not_none_indexes = {}\n        self.num_none_grad_updates = 0\n\n    def _init_aggregation_vars(self, grads):\n        \"\"\"\n        Initializes the counter that is used when to communicate and aggregate gradients\n        and the tensorflow variables that store the locally aggregated gradients.\n        \"\"\"\n        variable_scope_name = \"aggregation_variables_\" + str(self.rank)\n        with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):\n            self.counter = tf.compat.v1.get_variable(\n                \"aggregation_counter\", shape=(), dtype=tf.int32,\n                trainable=False, initializer=tf.compat.v1.zeros_initializer(),\n                collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES],\n            )\n            for idx, grad in enumerate(grads):\n                # Handle IndexedSlices.\n                if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                    grad = tf.convert_to_tensor(grad)\n                elif isinstance(grad, tf.IndexedSlices):\n                    raise ValueError(\n                        \"IndexedSlices are not supported when \"\n                        \"`backward_passes_per_step` > 1 and \"\n                        \"`sparse_as_dense` is False.\"\n                    )\n\n                # Handle grads that are None.\n                if grad is None:\n                    self.num_none_grad_updates += 1\n                    continue\n                self.not_none_indexes[idx] = len(self.locally_aggregated_grads)\n\n                # Create shadow variable.\n                grad_aggregation_variable_name = str(idx)\n                zero_grad = tf.zeros(shape=grad.get_shape().as_list(), dtype=grad.dtype)\n                grad_aggregation_variable = tf.compat.v1.get_variable(\n                    grad_aggregation_variable_name,\n                    trainable=False,\n                    initializer=zero_grad,\n                    collections=[\n                        tf.compat.v1.GraphKeys.LOCAL_VARIABLES,\n                        \"aggregating_collection\"],\n                )\n                self.locally_aggregated_grads.append(grad_aggregation_variable)\n            assert len(self.locally_aggregated_grads) + \\\n                self.num_none_grad_updates == len(grads)\n\n        # We expect to get a `sess` when we need to manually do a `sess.run(...)`\n        # for the variables to be initialized. This is the `tf.keras`\n        # optimizers.\n        if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:\n            session = tf.compat.v1.keras.backend.get_session(op_input_list=())\n            vars_init_op = tf.compat.v1.variables_initializer(\n                [self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]\n            )\n            session.run(vars_init_op)\n\n    def _clear_grads(self):\n        clear_ops_list = []\n        for idx, grad_aggregator in enumerate(self.locally_aggregated_grads):\n            clear_op = grad_aggregator.assign(grad_aggregator.initial_value)\n            clear_ops_list.append(clear_op)\n        return tf.group(*clear_ops_list)\n\n    def _aggregate_grads(self, grads):\n        aggregation_ops_list = []\n        grads = get_not_none_from_list(grads)\n        assert len(grads) == len(self.locally_aggregated_grads)\n\n        # Apply new gradient updates to the local copy.\n        for idx, grad in enumerate(grads):\n            if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                grad = tf.convert_to_tensor(grad)\n\n            updated_grad_aggregator = self.locally_aggregated_grads[idx].assign_add(\n                grad)\n            aggregation_ops_list.append(updated_grad_aggregator)\n\n        return aggregation_ops_list\n\n    def _allreduce_grads_helper(self, grads, vars):\n        # Read in latest variables values.\n        aggregated_grads = []\n        aggregation_read_ops_list = []\n        for idx, locally_aggregated_grad in enumerate(\n                self.locally_aggregated_grads):\n            aggregated_grads.append(locally_aggregated_grad.read_value())\n            aggregation_read_ops_list.append(aggregated_grads[idx])\n        aggregation_read_ops = tf.group(*aggregation_read_ops_list)\n\n        with tf.control_dependencies([aggregation_read_ops]):\n            averaged_gradients = self._allreduce_grads(aggregated_grads, vars)\n\n            # Reset counter.\n            with tf.control_dependencies([g.op for g in averaged_gradients if g is not None]):\n                reset_op = self.counter.assign(\n                    tf.constant(0), use_locking=True)\n\n            # Divide by backward_passes_per_step if\n            # average_aggregated_gradients is True.\n            with tf.control_dependencies([reset_op]):\n                gradient_divisor = self.backward_passes_per_step if \\\n                    self.average_aggregated_gradients else 1\n\n                averaged_gradients = apply_op_to_not_none_tensors(\n                    tf.divide,\n                    averaged_gradients,\n                    gradient_divisor,\n                )\n                return averaged_gradients\n\n    def compute_gradients(self, grads, vars):\n        \"\"\"\n        Applies the new gradient updates the locally aggregated gradients, and\n        performs cross-machine communication every backward_passes_per_step\n        times it is called.\n        \"\"\"\n        self._init_aggregation_vars(grads)\n\n        # Clear the locally aggregated gradients when the counter is at zero.\n        clear_op = tf.cond(\n            pred=tf.equal(self.counter, 0),\n            true_fn=lambda: self._clear_grads(),\n            false_fn=tf.no_op\n        )\n\n        # Add new gradients to the locally aggregated gradients.\n        with tf.control_dependencies([clear_op]):\n            aggregation_ops_list = self._aggregate_grads(grads)\n\n        # Increment the counter once new gradients have been applied.\n        aggregation_ops = tf.group(*aggregation_ops_list)\n        with tf.control_dependencies([aggregation_ops]):\n            update_counter = self.counter.assign_add(tf.constant(1))\n\n        with tf.control_dependencies([update_counter]):\n            grads = get_not_none_from_list(grads)\n            assert len(grads) == len(self.locally_aggregated_grads)\n\n            # Allreduce locally aggregated gradients when the counter is equivalent to\n            # `backward_passes_per_step`. This the condition is true, it also resets\n            # the counter back to 0.\n            allreduced_grads = tf.cond(\n                tf.equal(self.counter, self.backward_passes_per_step),\n                lambda: self._allreduce_grads_helper(grads, vars),\n                lambda: grads,\n            )\n\n            # Handle case where there is only one variable.\n            if not isinstance(allreduced_grads, (list, tuple)):\n                allreduced_grads = (allreduced_grads,)\n            assert len(allreduced_grads) == len(self.locally_aggregated_grads)\n\n            # Insert gradients that are None back in.\n            allreduced_grads = [\n                allreduced_grads[self.not_none_indexes[idx]] if idx in self.not_none_indexes else None\n                for idx in range(len(self.locally_aggregated_grads) + self.num_none_grad_updates)\n            ]\n            assert len(allreduced_grads) == len(\n                self.locally_aggregated_grads) + self.num_none_grad_updates\n\n        # If gradients have not been allreduced this batch, we return the gradients\n        # that were submitted as the updates (the input).\n        return allreduced_grads\n\n    def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):\n        \"\"\"\n        Apply updates every backward_passes_per_step, which lines up with\n        the batches on which we communicated the locally aggregated gradients.\n        \"\"\"\n        flattended_args0 = [item for tup in args[0] for item in tup]\n\n        # Since we skip applying updates when the counter is not at zero we\n        # still want to increment the global step if it is being tracked\n        # (e.g., Tensorflow Estimators).\n        def increment_global_step_counter():\n            global_step_counter = tf.compat.v1.train.get_global_step()\n            if global_step_counter is None:\n                return tf.no_op()\n            return global_step_counter.assign_add(\n                tf.constant(1, dtype=tf.int64),\n                use_locking=True,\n                read_value=False\n            )\n\n        # Increment global step on iterations where we don't call `apply_gradients()`.\n        cond_increment_global_step_counter = tf.cond(\n            pred=tf.equal(self.counter, 0),\n            true_fn=tf.no_op,\n            false_fn=increment_global_step_counter,\n        )\n        flattended_args0.append(cond_increment_global_step_counter)\n\n        # If optimizer tracks iterations, we increment it on steps where we\n        # are not going to call `apply_gradients()`.\n        def increment_optimizer_iteration():\n            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n                return optimizer._iterations.assign_add(1).op\n            return tf.no_op()\n\n        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):\n            return tf.cond(\n                pred=tf.equal(self.counter, 0),\n                true_fn=apply_grads_closure,\n                false_fn=increment_optimizer_iteration,\n            )\n", "path": "horovod/tensorflow/gradient_aggregation.py"}], "after_files": [{"content": "import tensorflow as tf\n\n\ndef apply_op_to_not_none_tensors(tensor_op, tensors, *args):\n    return [\n        tensor_op(\n            tensor,\n            *args\n        ) if tensor is not None else tensor for tensor in tensors]\n\n\ndef get_not_none_from_list(tensor_list):\n    return [x for x in tensor_list if x is not None]\n\n\nclass LocalGradientAggregationHelper:\n    \"\"\"\n    LocalGradientAggregationHelper aggregates gradient updates locally,\n    and communicates the updates across machines only once every\n    backward_passes_per_step. Only supports graph mode execution.\n    \"\"\"\n\n    _OPTIMIZER_TYPE_KERAS = \"optimizer_type_keras\"\n    _OPTIMIZER_TYPE_LEGACY = \"optimizer_type_legacy\"\n\n    def __init__(\n            self,\n            backward_passes_per_step,\n            allreduce_func,\n            sparse_as_dense,\n            average_aggregated_gradients,\n            rank,\n            optimizer_type):\n        self._allreduce_grads = allreduce_func\n\n        # backward_passes_per_step controls how often gradient updates are\n        # synchronized.\n        self.backward_passes_per_step = backward_passes_per_step\n        if self.backward_passes_per_step <= 0:\n            raise ValueError(\"backward_passes_per_step must be > 0\")\n\n        # average_aggregated_gradients controls whether gradient updates that are\n        # aggregated, should be divided by `backward_passes_per_step`.\n        self.average_aggregated_gradients = average_aggregated_gradients\n\n        # This is going to be [N] data structure holding the aggregated gradient updates\n        # N is the number of parameters.\n        self.locally_aggregated_grads = []\n\n        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`\n        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is\n        # equal to 0.\n        self.counter = None\n\n        self.sparse_as_dense = sparse_as_dense\n        self.rank = rank\n        self.optimizer_type = optimizer_type\n\n        # Contains the mapping of indexes of grad updates that are not None to their index in\n        # locally_aggregated_grads which only contains not None gradients. When performing\n        # gradient aggregation we have to remove them from the list of grads prior to passing\n        # the list into a tf.cond().\n        self.not_none_indexes = {}\n        self.num_none_grad_updates = 0\n\n    def _init_aggregation_vars(self, grads):\n        \"\"\"\n        Initializes the counter that is used when to communicate and aggregate gradients\n        and the tensorflow variables that store the locally aggregated gradients.\n        \"\"\"\n        variable_scope_name = \"aggregation_variables_\" + str(self.rank)\n        with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):\n            self.counter = tf.compat.v1.get_variable(\n                \"aggregation_counter\", shape=(), dtype=tf.int32,\n                trainable=False, initializer=tf.compat.v1.zeros_initializer(),\n                collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES],\n            )\n            for idx, grad in enumerate(grads):\n                # Handle IndexedSlices.\n                if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                    grad = tf.convert_to_tensor(grad)\n                elif isinstance(grad, tf.IndexedSlices):\n                    raise ValueError(\n                        \"IndexedSlices are not supported when \"\n                        \"`backward_passes_per_step` > 1 and \"\n                        \"`sparse_as_dense` is False.\"\n                    )\n\n                # Handle grads that are None.\n                if grad is None:\n                    self.num_none_grad_updates += 1\n                    continue\n                self.not_none_indexes[idx] = len(self.locally_aggregated_grads)\n\n                # Create shadow variable.\n                grad_aggregation_variable_name = str(idx)\n                zero_grad = tf.zeros(shape=grad.get_shape().as_list(), dtype=grad.dtype)\n                grad_aggregation_variable = tf.compat.v1.get_variable(\n                    grad_aggregation_variable_name,\n                    trainable=False,\n                    initializer=zero_grad,\n                    collections=[\n                        tf.compat.v1.GraphKeys.LOCAL_VARIABLES,\n                        \"aggregating_collection\"],\n                )\n                self.locally_aggregated_grads.append(grad_aggregation_variable)\n            assert len(self.locally_aggregated_grads) + \\\n                self.num_none_grad_updates == len(grads)\n\n        # We expect to get a `sess` when we need to manually do a `sess.run(...)`\n        # for the variables to be initialized. This is the `tf.keras`\n        # optimizers.\n        if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:\n            session = tf.compat.v1.keras.backend.get_session(op_input_list=())\n            vars_init_op = tf.compat.v1.variables_initializer(\n                [self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]\n            )\n            session.run(vars_init_op)\n\n    def _clear_grads(self):\n        clear_ops_list = []\n        for idx, grad_aggregator in enumerate(self.locally_aggregated_grads):\n            clear_op = grad_aggregator.assign(grad_aggregator.initial_value)\n            clear_ops_list.append(clear_op)\n        return tf.group(*clear_ops_list)\n\n    def _aggregate_grads(self, grads):\n        aggregation_ops_list = []\n        grads = get_not_none_from_list(grads)\n        assert len(grads) == len(self.locally_aggregated_grads)\n\n        # Apply new gradient updates to the local copy.\n        for idx, grad in enumerate(grads):\n            if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                grad = tf.convert_to_tensor(grad)\n\n            updated_grad_aggregator = self.locally_aggregated_grads[idx].assign_add(\n                grad)\n            aggregation_ops_list.append(updated_grad_aggregator)\n\n        return aggregation_ops_list\n\n    def _allreduce_grads_helper(self, grads, vars):\n        # Read in latest variables values.\n        aggregated_grads = []\n        aggregation_read_ops_list = []\n        for idx, locally_aggregated_grad in enumerate(\n                self.locally_aggregated_grads):\n            aggregated_grads.append(locally_aggregated_grad.read_value())\n            aggregation_read_ops_list.append(aggregated_grads[idx])\n        aggregation_read_ops = tf.group(*aggregation_read_ops_list)\n\n        with tf.control_dependencies([aggregation_read_ops]):\n            averaged_gradients = self._allreduce_grads(aggregated_grads, vars)\n\n            # Reset counter.\n            with tf.control_dependencies([g.op for g in averaged_gradients if g is not None]):\n                reset_op = self.counter.assign(\n                    tf.constant(0), use_locking=True)\n\n            # Divide by backward_passes_per_step if\n            # average_aggregated_gradients is True.\n            with tf.control_dependencies([reset_op]):\n                gradient_divisor = self.backward_passes_per_step if \\\n                    self.average_aggregated_gradients else 1\n\n                averaged_gradients = apply_op_to_not_none_tensors(\n                    tf.divide,\n                    averaged_gradients,\n                    gradient_divisor,\n                )\n                return averaged_gradients\n\n    def compute_gradients(self, grads, vars):\n        \"\"\"\n        Applies the new gradient updates the locally aggregated gradients, and\n        performs cross-machine communication every backward_passes_per_step\n        times it is called.\n        \"\"\"\n        self._init_aggregation_vars(grads)\n\n        # Clear the locally aggregated gradients when the counter is at zero.\n        clear_op = tf.cond(\n            pred=tf.equal(self.counter, 0),\n            true_fn=lambda: self._clear_grads(),\n            false_fn=tf.no_op\n        )\n\n        # Add new gradients to the locally aggregated gradients.\n        with tf.control_dependencies([clear_op]):\n            aggregation_ops_list = self._aggregate_grads(grads)\n\n        # Increment the counter once new gradients have been applied.\n        aggregation_ops = tf.group(*aggregation_ops_list)\n        with tf.control_dependencies([aggregation_ops]):\n            update_counter = self.counter.assign_add(tf.constant(1))\n\n        with tf.control_dependencies([update_counter]):\n            grads = get_not_none_from_list(grads)\n            assert len(grads) == len(self.locally_aggregated_grads)\n\n            # Allreduce locally aggregated gradients when the counter is equivalent to\n            # `backward_passes_per_step`. This the condition is true, it also resets\n            # the counter back to 0.\n            allreduced_grads = tf.cond(\n                tf.equal(self.counter, self.backward_passes_per_step),\n                lambda: self._allreduce_grads_helper(grads, vars),\n                lambda: grads,\n            )\n\n            # Handle case where there is only one variable.\n            if not isinstance(allreduced_grads, (list, tuple)):\n                allreduced_grads = (allreduced_grads,)\n            assert len(allreduced_grads) == len(self.locally_aggregated_grads)\n\n            # Insert gradients that are None back in.\n            allreduced_grads = [\n                allreduced_grads[self.not_none_indexes[idx]] if idx in self.not_none_indexes else None\n                for idx in range(len(self.locally_aggregated_grads) + self.num_none_grad_updates)\n            ]\n            assert len(allreduced_grads) == len(\n                self.locally_aggregated_grads) + self.num_none_grad_updates\n\n        # If gradients have not been allreduced this batch, we return the gradients\n        # that were submitted as the updates (the input).\n        return allreduced_grads\n\n    def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):\n        \"\"\"\n        Apply updates every backward_passes_per_step, which lines up with\n        the batches on which we communicated the locally aggregated gradients.\n        \"\"\"\n        flattended_args0 = [item for tup in args[0] for item in tup]\n\n        # If optimizer tracks iterations, we increment it on steps where we\n        # are not going to call `apply_gradients()`.\n        def increment_optimizer_iteration():\n            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n                return optimizer._iterations.assign_add(1).op\n            return tf.no_op()\n\n        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):\n            train_op = tf.cond(\n                pred=tf.equal(self.counter, 0),\n                true_fn=apply_grads_closure,\n                false_fn=increment_optimizer_iteration,\n            )\n\n        # Since we skip applying updates when the counter is not at zero we\n        # still want to increment the global step if it is being tracked\n        # (e.g., Tensorflow Estimators).\n        def increment_global_step_counter():\n            global_step_counter = tf.compat.v1.train.get_global_step()\n            if global_step_counter is None:\n                return tf.no_op()\n            return global_step_counter.assign_add(\n                tf.constant(1, dtype=tf.int64),\n                use_locking=True,\n                read_value=False\n            )\n\n        with tf.control_dependencies([train_op]):\n            # Increment global step on iterations where we don't call `apply_gradients()`.\n            return tf.cond(\n                pred=tf.equal(self.counter, 0),\n                true_fn=tf.no_op,\n                false_fn=increment_global_step_counter,\n            )\n", "path": "horovod/tensorflow/gradient_aggregation.py"}]}