We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error in computing gradients when using allgather
**Environment:**
1. Framework: TensorFlow
2. Framework version: 2.0
3. Horovod version:  0.18.2

I am trying to get the median of a tensor computed across all batches and all processes. However, I got an error TypeError: Expected int32, got None of type 'NoneType' instead.It seems that computing gradients does not work well with horovod's allgather operation. A simple illustration of what I would like to achieve is as follows:

>with tf.GradientTape() as tape: 
&ensp;&ensp;&ensp;&ensp;my_tensor = compute_my_tensor() 
&ensp;&ensp;&ensp;&ensp;gathered_my_tensor = hvd.allgather(my_tensor)  
&ensp;&ensp;&ensp;&ensp;median = get_median(gathered_my_tensor)
&ensp;&ensp;&ensp;&ensp;loss = get_loss(my_tensor, median, training=True)
tape = hvd.DistributedGradientTape(tape)
grads = tape.gradient(loss, trainable_variables)
optimizer.apply_gradients(zip(grads, trainable_variables))

BTW, when I use eager mode of tensorflow, there will be no error



--- 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/mpi_ops.py`
Content:
```
1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2 # Modifications copyright (C) 2019 Uber Technologies, Inc.
3 # Modifications copyright Microsoft
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #     http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 # =============================================================================
17 """Inter-process communication using MPI."""
18 
19 import re
20 import tensorflow as tf
21 from tensorflow.python.framework import load_library
22 from tensorflow.python.framework import ops
23 from tensorflow.python.platform import resource_loader
24 
25 from horovod.common.util import get_ext_suffix, get_average_backwards_compatibility_fun, gpu_available, \
26     num_rank_is_power_2
27 from horovod.common.basics import HorovodBasics as _HorovodBasics
28 from horovod.tensorflow.util import _executing_eagerly
29 
30 
31 def _load_library(name):
32     """Loads a .so file containing the specified operators.
33 
34     Args:
35       name: The name of the .so file to load.
36 
37     Raises:
38       NotFoundError if were not able to load .so file.
39     """
40     filename = resource_loader.get_path_to_datafile(name)
41     library = load_library.load_op_library(filename)
42     return library
43 
44 
45 MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
46 
47 _basics = _HorovodBasics(__file__, 'mpi_lib')
48 
49 # import basic methods
50 init = _basics.init
51 shutdown = _basics.shutdown
52 size = _basics.size
53 local_size = _basics.local_size
54 rank = _basics.rank
55 local_rank = _basics.local_rank
56 mpi_threads_supported = _basics.mpi_threads_supported
57 mpi_enabled = _basics.mpi_enabled
58 mpi_built = _basics.mpi_built
59 gloo_enabled = _basics.gloo_enabled
60 gloo_built = _basics.gloo_built
61 nccl_built = _basics.nccl_built
62 ddl_built = _basics.ddl_built
63 ccl_built = _basics.ccl_built
64 
65 # import reduction op values
66 Average = _basics.Average
67 Sum = _basics.Sum
68 Adasum = _basics.Adasum
69 
70 is_homogeneous = _basics.is_homogeneous
71 
72 handle_average_backwards_compatibility = get_average_backwards_compatibility_fun(_basics)
73 
74 check_num_rank_power_of_2 = num_rank_is_power_2
75 
76 
77 # This function will create a default device map which includes all visible devices.
78 # Please run this function in a subprocess
79 def _check_has_gpu():
80     import tensorflow as tf
81     return tf.test.is_gpu_available()
82 
83 
84 def _normalize_name(name):
85     """Normalizes operation name to TensorFlow rules."""
86     return re.sub('[^a-zA-Z0-9_]', '_', name)
87 
88 
89 def _allreduce(tensor, name=None, op=Sum):
90     """An op which reduces an input tensor over all the Horovod processes. The
91     default reduction is a sum.
92 
93     The reduction operation is keyed by the name of the op. The tensor type and
94     shape must be the same on all Horovod processes for a given name. The reduction
95     will not start until all processes are ready to send and receive the tensor.
96 
97     Returns:
98       A tensor of the same shape and type as `tensor`, summed across all
99       processes.
100     """
101     if name is None and not _executing_eagerly():
102         name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
103     return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op)
104 
105 
106 @ops.RegisterGradient('HorovodAllreduce')
107 def _allreduce_grad(op, grad):
108     """Gradient for allreduce op.
109 
110     Args:
111       op: An operation.
112       grad: `Tensor` gradient with respect to the output of the op.
113 
114     Returns:
115       The gradient with respect to the input of the op.
116     """
117     reduce_op = op.get_attr('reduce_op')
118     return _allreduce(grad, op=reduce_op)
119 
120 
121 def allgather(tensor, name=None):
122     """An op which concatenates the input tensor with the same input tensor on
123     all other Horovod processes.
124 
125     The concatenation is done on the first dimension, so the input tensors on the
126     different processes must have the same rank and shape, except for the first
127     dimension, which is allowed to be different.
128 
129     Returns:
130       A tensor of the same type as `tensor`, concatenated on dimension zero
131       across all processes. The shape is identical to the input shape, except for
132       the first dimension, which may be greater and is the sum of all first
133       dimensions of the tensors in different Horovod processes.
134     """
135     if name is None and not _executing_eagerly():
136         name = 'HorovodAllgather_%s' % _normalize_name(tensor.name)
137     return MPI_LIB.horovod_allgather(tensor, name=name)
138 
139 
140 @ops.RegisterGradient('HorovodAllgather')
141 def _allgather_grad(op, grad):
142     """Gradient for allgather op.
143 
144     Args:
145       op: An operation.
146       grad: `Tensor` gradient with respect to the output of the op.
147 
148     Returns:
149       The gradient with respect to the input of the op.
150     """
151     grad = _allreduce(grad)
152 
153     with tf.device('/cpu:0'):
154         # Keep the tensor of split sizes on CPU.
155         x = op.inputs[0]
156         d0 = x.get_shape().as_list()[0]
157         d = tf.convert_to_tensor([d0], dtype=tf.int32)
158 
159         s = size()
160         d = tf.reshape(allgather(d), [s])
161 
162     splits = tf.split(grad, num_or_size_splits=d, axis=0)
163     return splits[rank()]
164 
165 
166 def broadcast(tensor, root_rank, name=None):
167     """An op which broadcasts the input tensor on root rank to the same input tensor
168     on all other Horovod processes.
169 
170     The broadcast operation is keyed by the name of the op. The tensor type and
171     shape must be the same on all Horovod processes for a given name. The broadcast
172     will not start until all processes are ready to send and receive the tensor.
173 
174     Returns:
175       A tensor of the same shape and type as `tensor`, with the value broadcasted
176       from root rank.
177     """
178     if name is None and not _executing_eagerly():
179         name = 'HorovodBroadcast_%s' % _normalize_name(tensor.name)
180     return MPI_LIB.horovod_broadcast(tensor, name=name, root_rank=root_rank)
181 
182 
183 @ops.RegisterGradient('HorovodBroadcast')
184 def _broadcast_grad(op, grad):
185     """Gradient for broadcast op.
186 
187     Args:
188       op: An operation.
189       grad: `Tensor` gradient with respect to the output of the op.
190 
191     Returns:
192       The gradient with respect to the input of the op.
193     """
194     root_rank = op.get_attr('root_rank')
195     grad_reduced = _allreduce(grad)
196     if rank() != root_rank:
197         return grad_reduced * 0
198     return grad_reduced
199 
200 
201 def join():
202     return MPI_LIB.horovod_join()
203 
204 
205 def size_op(name=None):
206     """An op that returns the number of Horovod processes.
207 
208     This operation determines the return value at the graph execution time,
209     rather than at the graph construction time, and so allows for a graph to be
210     constructed in a different environment than where it will be executed.
211 
212     Returns:
213       An integer scalar containing the number of Horovod processes.
214     """
215     return MPI_LIB.horovod_size(name=name)
216 
217 
218 ops.NotDifferentiable('HorovodSize')
219 
220 
221 def local_size_op(name=None):
222     """An op that returns the number of Horovod processes within the
223     node the current process is running on.
224 
225     This operation determines the return value at the graph execution time,
226     rather than at the graph construction time, and so allows for a graph to be
227     constructed in a different environment than where it will be executed.
228 
229     Returns:
230       An integer scalar containing the number of local Horovod processes.
231     """
232     return MPI_LIB.horovod_local_size(name=name)
233 
234 
235 ops.NotDifferentiable('HorovodLocalSize')
236 
237 
238 def rank_op(name=None):
239     """An op that returns the Horovod rank of the calling process.
240 
241     This operation determines the return value at the graph execution time,
242     rather than at the graph construction time, and so allows for a graph to be
243     constructed in a different environment than where it will be executed.
244 
245     Returns:
246       An integer scalar with the Horovod rank of the calling process.
247     """
248     return MPI_LIB.horovod_rank(name=name)
249 
250 
251 ops.NotDifferentiable('HorovodRank')
252 
253 
254 def local_rank_op(name=None):
255     """An op that returns the local Horovod rank of the calling process, within the
256     node that it is running on. For example, if there are seven processes running
257     on a node, their local ranks will be zero through six, inclusive.
258 
259     This operation determines the return value at the graph execution time,
260     rather than at the graph construction time, and so allows for a graph to be
261     constructed in a different environment than where it will be executed.
262 
263     Returns:
264       An integer scalar with the local Horovod rank of the calling process.
265     """
266     return MPI_LIB.horovod_rank(name=name)
267 
268 
269 ops.NotDifferentiable('HorovodLocalRank')
270
```
--- 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/mpi_ops.py b/horovod/tensorflow/mpi_ops.py
--- a/horovod/tensorflow/mpi_ops.py
+++ b/horovod/tensorflow/mpi_ops.py
@@ -152,8 +152,8 @@
     with tf.device('/cpu:0'):
         # Keep the tensor of split sizes on CPU.
         x = op.inputs[0]
-        d0 = x.get_shape().as_list()[0]
-        d = tf.convert_to_tensor([d0], dtype=tf.int32)
+        d = tf.shape(x)
+        d = tf.reshape(d[0], [1])
 
         s = size()
         d = tf.reshape(allgather(d), [s])

{"golden_diff": "diff --git a/horovod/tensorflow/mpi_ops.py b/horovod/tensorflow/mpi_ops.py\n--- a/horovod/tensorflow/mpi_ops.py\n+++ b/horovod/tensorflow/mpi_ops.py\n@@ -152,8 +152,8 @@\n     with tf.device('/cpu:0'):\n         # Keep the tensor of split sizes on CPU.\n         x = op.inputs[0]\n-        d0 = x.get_shape().as_list()[0]\n-        d = tf.convert_to_tensor([d0], dtype=tf.int32)\n+        d = tf.shape(x)\n+        d = tf.reshape(d[0], [1])\n \n         s = size()\n         d = tf.reshape(allgather(d), [s])\n", "issue": "Error in computing gradients when using allgather\n**Environment:**\r\n1. Framework: TensorFlow\r\n2. Framework version: 2.0\r\n3. Horovod version:  0.18.2\r\n\r\nI am trying to get the median of a tensor computed across all batches and all processes. However, I got an error TypeError: Expected int32, got None of type 'NoneType' instead.It seems that computing gradients does not work well with horovod's allgather operation. A simple illustration of what I would like to achieve is as follows:\r\n\r\n>with tf.GradientTape() as tape: \r\n&ensp;&ensp;&ensp;&ensp;my_tensor = compute_my_tensor() \r\n&ensp;&ensp;&ensp;&ensp;gathered_my_tensor = hvd.allgather(my_tensor)  \r\n&ensp;&ensp;&ensp;&ensp;median = get_median(gathered_my_tensor)\r\n&ensp;&ensp;&ensp;&ensp;loss = get_loss(my_tensor, median, training=True)\r\ntape = hvd.DistributedGradientTape(tape)\r\ngrads = tape.gradient(loss, trainable_variables)\r\noptimizer.apply_gradients(zip(grads, trainable_variables))\r\n\r\nBTW, when I use eager mode of tensorflow, there will be no error\r\n\r\n\n", "before_files": [{"content": "# Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n# Modifications copyright (C) 2019 Uber Technologies, Inc.\n# Modifications copyright Microsoft\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# =============================================================================\n\"\"\"Inter-process communication using MPI.\"\"\"\n\nimport re\nimport tensorflow as tf\nfrom tensorflow.python.framework import load_library\nfrom tensorflow.python.framework import ops\nfrom tensorflow.python.platform import resource_loader\n\nfrom horovod.common.util import get_ext_suffix, get_average_backwards_compatibility_fun, gpu_available, \\\n    num_rank_is_power_2\nfrom horovod.common.basics import HorovodBasics as _HorovodBasics\nfrom horovod.tensorflow.util import _executing_eagerly\n\n\ndef _load_library(name):\n    \"\"\"Loads a .so file containing the specified operators.\n\n    Args:\n      name: The name of the .so file to load.\n\n    Raises:\n      NotFoundError if were not able to load .so file.\n    \"\"\"\n    filename = resource_loader.get_path_to_datafile(name)\n    library = load_library.load_op_library(filename)\n    return library\n\n\nMPI_LIB = _load_library('mpi_lib' + get_ext_suffix())\n\n_basics = _HorovodBasics(__file__, 'mpi_lib')\n\n# import basic methods\ninit = _basics.init\nshutdown = _basics.shutdown\nsize = _basics.size\nlocal_size = _basics.local_size\nrank = _basics.rank\nlocal_rank = _basics.local_rank\nmpi_threads_supported = _basics.mpi_threads_supported\nmpi_enabled = _basics.mpi_enabled\nmpi_built = _basics.mpi_built\ngloo_enabled = _basics.gloo_enabled\ngloo_built = _basics.gloo_built\nnccl_built = _basics.nccl_built\nddl_built = _basics.ddl_built\nccl_built = _basics.ccl_built\n\n# import reduction op values\nAverage = _basics.Average\nSum = _basics.Sum\nAdasum = _basics.Adasum\n\nis_homogeneous = _basics.is_homogeneous\n\nhandle_average_backwards_compatibility = get_average_backwards_compatibility_fun(_basics)\n\ncheck_num_rank_power_of_2 = num_rank_is_power_2\n\n\n# This function will create a default device map which includes all visible devices.\n# Please run this function in a subprocess\ndef _check_has_gpu():\n    import tensorflow as tf\n    return tf.test.is_gpu_available()\n\n\ndef _normalize_name(name):\n    \"\"\"Normalizes operation name to TensorFlow rules.\"\"\"\n    return re.sub('[^a-zA-Z0-9_]', '_', name)\n\n\ndef _allreduce(tensor, name=None, op=Sum):\n    \"\"\"An op which reduces an input tensor over all the Horovod processes. The\n    default reduction is a sum.\n\n    The reduction operation is keyed by the name of the op. The tensor type and\n    shape must be the same on all Horovod processes for a given name. The reduction\n    will not start until all processes are ready to send and receive the tensor.\n\n    Returns:\n      A tensor of the same shape and type as `tensor`, summed across all\n      processes.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op)\n\n\n@ops.RegisterGradient('HorovodAllreduce')\ndef _allreduce_grad(op, grad):\n    \"\"\"Gradient for allreduce op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    reduce_op = op.get_attr('reduce_op')\n    return _allreduce(grad, op=reduce_op)\n\n\ndef allgather(tensor, name=None):\n    \"\"\"An op which concatenates the input tensor with the same input tensor on\n    all other Horovod processes.\n\n    The concatenation is done on the first dimension, so the input tensors on the\n    different processes must have the same rank and shape, except for the first\n    dimension, which is allowed to be different.\n\n    Returns:\n      A tensor of the same type as `tensor`, concatenated on dimension zero\n      across all processes. The shape is identical to the input shape, except for\n      the first dimension, which may be greater and is the sum of all first\n      dimensions of the tensors in different Horovod processes.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodAllgather_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_allgather(tensor, name=name)\n\n\n@ops.RegisterGradient('HorovodAllgather')\ndef _allgather_grad(op, grad):\n    \"\"\"Gradient for allgather op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    grad = _allreduce(grad)\n\n    with tf.device('/cpu:0'):\n        # Keep the tensor of split sizes on CPU.\n        x = op.inputs[0]\n        d0 = x.get_shape().as_list()[0]\n        d = tf.convert_to_tensor([d0], dtype=tf.int32)\n\n        s = size()\n        d = tf.reshape(allgather(d), [s])\n\n    splits = tf.split(grad, num_or_size_splits=d, axis=0)\n    return splits[rank()]\n\n\ndef broadcast(tensor, root_rank, name=None):\n    \"\"\"An op which broadcasts the input tensor on root rank to the same input tensor\n    on all other Horovod processes.\n\n    The broadcast operation is keyed by the name of the op. The tensor type and\n    shape must be the same on all Horovod processes for a given name. The broadcast\n    will not start until all processes are ready to send and receive the tensor.\n\n    Returns:\n      A tensor of the same shape and type as `tensor`, with the value broadcasted\n      from root rank.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodBroadcast_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_broadcast(tensor, name=name, root_rank=root_rank)\n\n\n@ops.RegisterGradient('HorovodBroadcast')\ndef _broadcast_grad(op, grad):\n    \"\"\"Gradient for broadcast op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    root_rank = op.get_attr('root_rank')\n    grad_reduced = _allreduce(grad)\n    if rank() != root_rank:\n        return grad_reduced * 0\n    return grad_reduced\n\n\ndef join():\n    return MPI_LIB.horovod_join()\n\n\ndef size_op(name=None):\n    \"\"\"An op that returns the number of Horovod processes.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar containing the number of Horovod processes.\n    \"\"\"\n    return MPI_LIB.horovod_size(name=name)\n\n\nops.NotDifferentiable('HorovodSize')\n\n\ndef local_size_op(name=None):\n    \"\"\"An op that returns the number of Horovod processes within the\n    node the current process is running on.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar containing the number of local Horovod processes.\n    \"\"\"\n    return MPI_LIB.horovod_local_size(name=name)\n\n\nops.NotDifferentiable('HorovodLocalSize')\n\n\ndef rank_op(name=None):\n    \"\"\"An op that returns the Horovod rank of the calling process.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar with the Horovod rank of the calling process.\n    \"\"\"\n    return MPI_LIB.horovod_rank(name=name)\n\n\nops.NotDifferentiable('HorovodRank')\n\n\ndef local_rank_op(name=None):\n    \"\"\"An op that returns the local Horovod rank of the calling process, within the\n    node that it is running on. For example, if there are seven processes running\n    on a node, their local ranks will be zero through six, inclusive.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar with the local Horovod rank of the calling process.\n    \"\"\"\n    return MPI_LIB.horovod_rank(name=name)\n\n\nops.NotDifferentiable('HorovodLocalRank')\n", "path": "horovod/tensorflow/mpi_ops.py"}], "after_files": [{"content": "# Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n# Modifications copyright (C) 2019 Uber Technologies, Inc.\n# Modifications copyright Microsoft\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# =============================================================================\n\"\"\"Inter-process communication using MPI.\"\"\"\n\nimport re\nimport tensorflow as tf\nfrom tensorflow.python.framework import load_library\nfrom tensorflow.python.framework import ops\nfrom tensorflow.python.platform import resource_loader\n\nfrom horovod.common.util import get_ext_suffix, get_average_backwards_compatibility_fun, gpu_available, \\\n    num_rank_is_power_2\nfrom horovod.common.basics import HorovodBasics as _HorovodBasics\nfrom horovod.tensorflow.util import _executing_eagerly\n\n\ndef _load_library(name):\n    \"\"\"Loads a .so file containing the specified operators.\n\n    Args:\n      name: The name of the .so file to load.\n\n    Raises:\n      NotFoundError if were not able to load .so file.\n    \"\"\"\n    filename = resource_loader.get_path_to_datafile(name)\n    library = load_library.load_op_library(filename)\n    return library\n\n\nMPI_LIB = _load_library('mpi_lib' + get_ext_suffix())\n\n_basics = _HorovodBasics(__file__, 'mpi_lib')\n\n# import basic methods\ninit = _basics.init\nshutdown = _basics.shutdown\nsize = _basics.size\nlocal_size = _basics.local_size\nrank = _basics.rank\nlocal_rank = _basics.local_rank\nmpi_threads_supported = _basics.mpi_threads_supported\nmpi_enabled = _basics.mpi_enabled\nmpi_built = _basics.mpi_built\ngloo_enabled = _basics.gloo_enabled\ngloo_built = _basics.gloo_built\nnccl_built = _basics.nccl_built\nddl_built = _basics.ddl_built\nccl_built = _basics.ccl_built\n\n# import reduction op values\nAverage = _basics.Average\nSum = _basics.Sum\nAdasum = _basics.Adasum\n\nis_homogeneous = _basics.is_homogeneous\n\nhandle_average_backwards_compatibility = get_average_backwards_compatibility_fun(_basics)\n\ncheck_num_rank_power_of_2 = num_rank_is_power_2\n\n\n# This function will create a default device map which includes all visible devices.\n# Please run this function in a subprocess\ndef _check_has_gpu():\n    import tensorflow as tf\n    return tf.test.is_gpu_available()\n\n\ndef _normalize_name(name):\n    \"\"\"Normalizes operation name to TensorFlow rules.\"\"\"\n    return re.sub('[^a-zA-Z0-9_]', '_', name)\n\n\ndef _allreduce(tensor, name=None, op=Sum):\n    \"\"\"An op which reduces an input tensor over all the Horovod processes. The\n    default reduction is a sum.\n\n    The reduction operation is keyed by the name of the op. The tensor type and\n    shape must be the same on all Horovod processes for a given name. The reduction\n    will not start until all processes are ready to send and receive the tensor.\n\n    Returns:\n      A tensor of the same shape and type as `tensor`, summed across all\n      processes.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op)\n\n\n@ops.RegisterGradient('HorovodAllreduce')\ndef _allreduce_grad(op, grad):\n    \"\"\"Gradient for allreduce op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    return _allreduce(grad)\n\n\ndef allgather(tensor, name=None):\n    \"\"\"An op which concatenates the input tensor with the same input tensor on\n    all other Horovod processes.\n\n    The concatenation is done on the first dimension, so the input tensors on the\n    different processes must have the same rank and shape, except for the first\n    dimension, which is allowed to be different.\n\n    Returns:\n      A tensor of the same type as `tensor`, concatenated on dimension zero\n      across all processes. The shape is identical to the input shape, except for\n      the first dimension, which may be greater and is the sum of all first\n      dimensions of the tensors in different Horovod processes.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodAllgather_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_allgather(tensor, name=name)\n\n\n@ops.RegisterGradient('HorovodAllgather')\ndef _allgather_grad(op, grad):\n    \"\"\"Gradient for allgather op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    grad = _allreduce(grad)\n\n    with tf.device('/cpu:0'):\n        # Keep the tensor of split sizes on CPU.\n        x = op.inputs[0]\n        d = tf.shape(x)\n        d = tf.reshape(d[0], [1])\n\n        s = size()\n        d = tf.reshape(allgather(d), [s])\n\n    splits = tf.split(grad, num_or_size_splits=d, axis=0)\n    return splits[rank()]\n\n\ndef broadcast(tensor, root_rank, name=None):\n    \"\"\"An op which broadcasts the input tensor on root rank to the same input tensor\n    on all other Horovod processes.\n\n    The broadcast operation is keyed by the name of the op. The tensor type and\n    shape must be the same on all Horovod processes for a given name. The broadcast\n    will not start until all processes are ready to send and receive the tensor.\n\n    Returns:\n      A tensor of the same shape and type as `tensor`, with the value broadcasted\n      from root rank.\n    \"\"\"\n    if name is None and not _executing_eagerly():\n        name = 'HorovodBroadcast_%s' % _normalize_name(tensor.name)\n    return MPI_LIB.horovod_broadcast(tensor, name=name, root_rank=root_rank)\n\n\n@ops.RegisterGradient('HorovodBroadcast')\ndef _broadcast_grad(op, grad):\n    \"\"\"Gradient for broadcast op.\n\n    Args:\n      op: An operation.\n      grad: `Tensor` gradient with respect to the output of the op.\n\n    Returns:\n      The gradient with respect to the input of the op.\n    \"\"\"\n    root_rank = op.get_attr('root_rank')\n    grad_reduced = _allreduce(grad)\n    if rank() != root_rank:\n        return grad_reduced * 0\n    return grad_reduced\n\n\ndef join():\n    return MPI_LIB.horovod_join()\n\n\ndef size_op(name=None):\n    \"\"\"An op that returns the number of Horovod processes.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar containing the number of Horovod processes.\n    \"\"\"\n    return MPI_LIB.horovod_size(name=name)\n\n\nops.NotDifferentiable('HorovodSize')\n\n\ndef local_size_op(name=None):\n    \"\"\"An op that returns the number of Horovod processes within the\n    node the current process is running on.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar containing the number of local Horovod processes.\n    \"\"\"\n    return MPI_LIB.horovod_local_size(name=name)\n\n\nops.NotDifferentiable('HorovodLocalSize')\n\n\ndef rank_op(name=None):\n    \"\"\"An op that returns the Horovod rank of the calling process.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar with the Horovod rank of the calling process.\n    \"\"\"\n    return MPI_LIB.horovod_rank(name=name)\n\n\nops.NotDifferentiable('HorovodRank')\n\n\ndef local_rank_op(name=None):\n    \"\"\"An op that returns the local Horovod rank of the calling process, within the\n    node that it is running on. For example, if there are seven processes running\n    on a node, their local ranks will be zero through six, inclusive.\n\n    This operation determines the return value at the graph execution time,\n    rather than at the graph construction time, and so allows for a graph to be\n    constructed in a different environment than where it will be executed.\n\n    Returns:\n      An integer scalar with the local Horovod rank of the calling process.\n    \"\"\"\n    return MPI_LIB.horovod_rank(name=name)\n\n\nops.NotDifferentiable('HorovodLocalRank')\n", "path": "horovod/tensorflow/mpi_ops.py"}]}