We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DistributedOptimzer is not compatible with keras.Optimizer
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet)
2. Framework version: 1.14
3. Horovod version: 0.16.4
7. Python version: 3.6.8

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?

**Bug report:**
Please describe errorneous behavior you're observing and steps to reproduce it.
Horovod DistributedOptimzer wrapper is not compatible with keras:



```
import tensorflow as tf
import horovod.tensorflow.keras as hvd

hvd.init()
opt = tf.keras.optimizers.Adam()
hopt = hvd.DistributedOptimizer(opt)
opt.get_config()
cfg = hopt.get_config()
opt_copy = opt.from_config(cfg)
opt_copy = opt.__class__.from_config(cfg)
hopt_copy = hopt.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'
hopt_copy = hopt.__class__.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'
```


--- 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/common/util.py`
Content:
```
1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2 # Modifications copyright (C) 2019 Uber Technologies, Inc.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 # =============================================================================
16 
17 from contextlib import contextmanager
18 import importlib
19 from multiprocessing import Process, Queue
20 import os
21 import sysconfig
22 
23 EXTENSIONS = ['tensorflow', 'torch', 'mxnet']
24 
25 
26 def get_ext_suffix():
27     """Determine library extension for various versions of Python."""
28     ext_suffix = sysconfig.get_config_var('EXT_SUFFIX')
29     if ext_suffix:
30         return ext_suffix
31 
32     ext_suffix = sysconfig.get_config_var('SO')
33     if ext_suffix:
34         return ext_suffix
35 
36     return '.so'
37 
38 
39 def get_extension_full_path(pkg_path, *args):
40     assert len(args) >= 1
41     dir_path = os.path.join(os.path.dirname(pkg_path), *args[:-1])
42     full_path = os.path.join(dir_path, args[-1] + get_ext_suffix())
43     return full_path
44 
45 
46 def check_extension(ext_name, ext_env_var, pkg_path, *args):
47     full_path = get_extension_full_path(pkg_path, *args)
48     if not os.path.exists(full_path):
49         raise ImportError(
50             'Extension %s has not been built.  If this is not expected, reinstall '
51             'Horovod with %s=1 to debug the build error.' % (ext_name, ext_env_var))
52 
53 
54 def _check_extension_lambda(ext_base_name, fn, fn_desc, verbose):
55     """
56     Tries to load the extension in a new process.  If successful, puts fn(ext)
57     to the queue or False otherwise.  Mutes all stdout/stderr.
58     """
59     def _target_fn(ext_base_name, fn, fn_desc, queue, verbose):
60         import importlib
61         import sys
62         import traceback
63 
64         if verbose:
65             print('Checking whether extension {ext_base_name} was {fn_desc}.'.format(
66                 ext_base_name=ext_base_name, fn_desc=fn_desc))
67         else:
68             # Suppress output
69             sys.stdout = open(os.devnull, 'w')
70             sys.stderr = open(os.devnull, 'w')
71 
72         try:
73             ext = importlib.import_module('.' + ext_base_name, 'horovod')
74             result = fn(ext)
75         except:
76             traceback.print_exc()
77             result = None
78 
79         if verbose:
80             print('Extension {ext_base_name} {flag} {fn_desc}.'.format(
81                 ext_base_name=ext_base_name, flag=('was' if result else 'was NOT'),
82                 fn_desc=fn_desc))
83 
84         queue.put(result)
85 
86     queue = Queue()
87     p = Process(target=_target_fn,
88                 args=(ext_base_name, fn, fn_desc, queue, verbose))
89     p.daemon = True
90     p.start()
91     p.join()
92     return queue.get_nowait()
93 
94 
95 def extension_available(ext_base_name, verbose=False):
96     available_fn = lambda ext: ext is not None
97     return _check_extension_lambda(
98         ext_base_name, available_fn, 'built', verbose) or False
99 
100 
101 def mpi_built(verbose=False):
102     for ext_base_name in EXTENSIONS:
103         built_fn = lambda ext: ext.mpi_built()
104         result = _check_extension_lambda(
105             ext_base_name, built_fn, 'built with MPI', verbose)
106         if result is not None:
107             return result
108     return False
109 
110 
111 def gloo_built(verbose=False):
112     for ext_base_name in EXTENSIONS:
113         built_fn = lambda ext: ext.gloo_built()
114         result = _check_extension_lambda(
115             ext_base_name, built_fn, 'built with Gloo', verbose)
116         if result is not None:
117             return result
118     return False
119 
120 
121 def nccl_built(verbose=False):
122     for ext_base_name in EXTENSIONS:
123         built_fn = lambda ext: ext.nccl_built()
124         result = _check_extension_lambda(
125             ext_base_name, built_fn, 'built with NCCL', verbose)
126         if result is not None:
127             return result
128     return False
129 
130 
131 def ddl_built(verbose=False):
132     for ext_base_name in EXTENSIONS:
133         built_fn = lambda ext: ext.ddl_built()
134         result = _check_extension_lambda(
135             ext_base_name, built_fn, 'built with DDL', verbose)
136         if result is not None:
137             return result
138     return False
139 
140 
141 def mlsl_built(verbose=False):
142     for ext_base_name in EXTENSIONS:
143         built_fn = lambda ext: ext.mlsl_built()
144         result = _check_extension_lambda(
145             ext_base_name, built_fn, 'built with MLSL', verbose)
146         if result is not None:
147             return result
148     return False
149 
150 
151 @contextmanager
152 def env(**kwargs):
153     # ignore args with None values
154     for k in list(kwargs.keys()):
155         if kwargs[k] is None:
156             del kwargs[k]
157 
158     # backup environment
159     backup = {}
160     for k in kwargs.keys():
161         backup[k] = os.environ.get(k)
162 
163     # set new values & yield
164     for k, v in kwargs.items():
165         os.environ[k] = v
166 
167     try:
168         yield
169     finally:
170         # restore environment
171         for k in kwargs.keys():
172             if backup[k] is not None:
173                 os.environ[k] = backup[k]
174             else:
175                 del os.environ[k]
176
```

Path: `horovod/_keras/__init__.py`
Content:
```
1 # Copyright 2017 Uber Technologies, Inc. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 import horovod.tensorflow as hvd
17 import tensorflow as tf
18 
19 
20 def create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,
21                                  compression, sparse_as_dense):
22     class _DistributedOptimizer(keras.optimizers.Optimizer):
23         def __init__(self, name, device_dense, device_sparse, compression, sparse_as_dense,
24                      config):
25             if name is None:
26                 name = "Distributed%s" % self.__class__.__base__.__name__
27             self._name = name
28             self._device_dense = device_dense
29             self._device_sparse = device_sparse
30             self._compression = compression
31             self._sparse_as_dense = sparse_as_dense
32             self._get_gradients_used = False
33             super(self.__class__, self).__init__(**config)
34 
35         def get_gradients(self, loss, params):
36             """
37             Compute gradients of all trainable variables.
38 
39             See Optimizer.get_gradients() for more info.
40 
41             In DistributedOptimizer, get_gradients() is overriden to also
42             allreduce the gradients before returning them.
43             """
44             self._get_gradients_used = True
45             gradients = super(self.__class__, self).get_gradients(loss, params)
46             if hvd.size() > 1:
47                 averaged_gradients = []
48                 with tf.name_scope(self._name + "_Allreduce"):
49                     for grad in gradients:
50                         if grad is not None:
51                             if self._sparse_as_dense and \
52                                     isinstance(grad, tf.IndexedSlices):
53                                 grad = tf.convert_to_tensor(grad)
54                             avg_grad = hvd.allreduce(grad,
55                                                      device_dense=self._device_dense,
56                                                      device_sparse=self._device_sparse,
57                                                      compression=self._compression)
58                             averaged_gradients.append(avg_grad)
59                         else:
60                             averaged_gradients.append(None)
61                     return averaged_gradients
62             else:
63                 return gradients
64 
65         def apply_gradients(self, *args, **kwargs):
66             if not self._get_gradients_used:
67                   raise Exception('`apply_gradients()` was called without a call to '
68                                   '`get_gradients()`. If you\'re using TensorFlow 2.0, '
69                                   'please specify `experimental_run_tf_function=False` in '
70                                   '`compile()`.')
71             return super(self.__class__, self).apply_gradients(*args, **kwargs)
72 
73     # We dynamically create a new class that inherits from the optimizer that was passed in.
74     # The goal is to override get_gradients() method with an allreduce implementation.
75     # This class will have the same name as the optimizer it's wrapping, so that the saved
76     # model could be easily restored without Horovod.
77     cls = type(optimizer.__class__.__name__, (optimizer.__class__,),
78                dict(_DistributedOptimizer.__dict__))
79     return cls(name, device_dense, device_sparse, compression, sparse_as_dense,
80                optimizer.get_config())
81 
82 
83 def _eval(backend, op_or_result):
84     if hvd._executing_eagerly():
85         return op_or_result
86     else:
87         return backend.get_session().run(op_or_result)
88 
89 
90 if hasattr(hvd, 'broadcast_global_variables'):
91     def broadcast_global_variables(backend, root_rank):
92         return _eval(backend, hvd.broadcast_global_variables(root_rank))
93 
94 
95 def allreduce(backend, value, name, average):
96     return _eval(backend, hvd.allreduce(tf.constant(value, name=name), average=average))
97 
98 
99 def allgather(backend, value, name):
100     return _eval(backend, hvd.allgather(tf.constant(value, name=name)))
101 
102 
103 def broadcast(backend, value, root_rank, name):
104     return _eval(backend, hvd.broadcast(tf.constant(value, name=name), root_rank))
105 
106 
107 def load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects):
108     horovod_objects = {
109         subclass.__name__.lower(): wrap_optimizer(subclass)
110         for subclass in keras.optimizers.Optimizer.__subclasses__()
111         if subclass.__module__ == keras.optimizers.Optimizer.__module__
112     }
113 
114     if custom_optimizers is not None:
115         horovod_objects.update({
116             cls.__name__: wrap_optimizer(cls)
117             for cls in custom_optimizers
118         })
119 
120     if custom_objects is not None:
121         horovod_objects.update(custom_objects)
122 
123     return keras.models.load_model(filepath, custom_objects=horovod_objects)
124
```
--- 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/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -70,6 +70,10 @@
                                   '`compile()`.')
             return super(self.__class__, self).apply_gradients(*args, **kwargs)
 
+        @classmethod
+        def from_config(cls, cfg):
+            return cls(name, device_dense, device_sparse, compression, sparse_as_dense, cfg)
+
     # We dynamically create a new class that inherits from the optimizer that was passed in.
     # The goal is to override get_gradients() method with an allreduce implementation.
     # This class will have the same name as the optimizer it's wrapping, so that the saved
diff --git a/horovod/common/util.py b/horovod/common/util.py
--- a/horovod/common/util.py
+++ b/horovod/common/util.py
@@ -15,7 +15,6 @@
 # =============================================================================
 
 from contextlib import contextmanager
-import importlib
 from multiprocessing import Process, Queue
 import os
 import sysconfig

{"golden_diff": "diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py\n--- a/horovod/_keras/__init__.py\n+++ b/horovod/_keras/__init__.py\n@@ -70,6 +70,10 @@\n                                   '`compile()`.')\n             return super(self.__class__, self).apply_gradients(*args, **kwargs)\n \n+        @classmethod\n+        def from_config(cls, cfg):\n+            return cls(name, device_dense, device_sparse, compression, sparse_as_dense, cfg)\n+\n     # We dynamically create a new class that inherits from the optimizer that was passed in.\n     # The goal is to override get_gradients() method with an allreduce implementation.\n     # This class will have the same name as the optimizer it's wrapping, so that the saved\ndiff --git a/horovod/common/util.py b/horovod/common/util.py\n--- a/horovod/common/util.py\n+++ b/horovod/common/util.py\n@@ -15,7 +15,6 @@\n # =============================================================================\n \n from contextlib import contextmanager\n-import importlib\n from multiprocessing import Process, Queue\n import os\n import sysconfig\n", "issue": "DistributedOptimzer is not compatible with keras.Optimizer\n**Environment:**\r\n1. Framework: (TensorFlow, Keras, PyTorch, MXNet)\r\n2. Framework version: 1.14\r\n3. Horovod version: 0.16.4\r\n7. Python version: 3.6.8\r\n\r\n**Checklist:**\r\n1. Did you search issues to find if somebody asked this question before?\r\n2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?\r\n3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?\r\n4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?\r\n\r\n**Bug report:**\r\nPlease describe errorneous behavior you're observing and steps to reproduce it.\r\nHorovod DistributedOptimzer wrapper is not compatible with keras:\r\n\r\n\r\n\r\n```\r\nimport tensorflow as tf\r\nimport horovod.tensorflow.keras as hvd\r\n\r\nhvd.init()\r\nopt = tf.keras.optimizers.Adam()\r\nhopt = hvd.DistributedOptimizer(opt)\r\nopt.get_config()\r\ncfg = hopt.get_config()\r\nopt_copy = opt.from_config(cfg)\r\nopt_copy = opt.__class__.from_config(cfg)\r\nhopt_copy = hopt.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'\r\nhopt_copy = hopt.__class__.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'\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#\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\nfrom contextlib import contextmanager\nimport importlib\nfrom multiprocessing import Process, Queue\nimport os\nimport sysconfig\n\nEXTENSIONS = ['tensorflow', 'torch', 'mxnet']\n\n\ndef get_ext_suffix():\n    \"\"\"Determine library extension for various versions of Python.\"\"\"\n    ext_suffix = sysconfig.get_config_var('EXT_SUFFIX')\n    if ext_suffix:\n        return ext_suffix\n\n    ext_suffix = sysconfig.get_config_var('SO')\n    if ext_suffix:\n        return ext_suffix\n\n    return '.so'\n\n\ndef get_extension_full_path(pkg_path, *args):\n    assert len(args) >= 1\n    dir_path = os.path.join(os.path.dirname(pkg_path), *args[:-1])\n    full_path = os.path.join(dir_path, args[-1] + get_ext_suffix())\n    return full_path\n\n\ndef check_extension(ext_name, ext_env_var, pkg_path, *args):\n    full_path = get_extension_full_path(pkg_path, *args)\n    if not os.path.exists(full_path):\n        raise ImportError(\n            'Extension %s has not been built.  If this is not expected, reinstall '\n            'Horovod with %s=1 to debug the build error.' % (ext_name, ext_env_var))\n\n\ndef _check_extension_lambda(ext_base_name, fn, fn_desc, verbose):\n    \"\"\"\n    Tries to load the extension in a new process.  If successful, puts fn(ext)\n    to the queue or False otherwise.  Mutes all stdout/stderr.\n    \"\"\"\n    def _target_fn(ext_base_name, fn, fn_desc, queue, verbose):\n        import importlib\n        import sys\n        import traceback\n\n        if verbose:\n            print('Checking whether extension {ext_base_name} was {fn_desc}.'.format(\n                ext_base_name=ext_base_name, fn_desc=fn_desc))\n        else:\n            # Suppress output\n            sys.stdout = open(os.devnull, 'w')\n            sys.stderr = open(os.devnull, 'w')\n\n        try:\n            ext = importlib.import_module('.' + ext_base_name, 'horovod')\n            result = fn(ext)\n        except:\n            traceback.print_exc()\n            result = None\n\n        if verbose:\n            print('Extension {ext_base_name} {flag} {fn_desc}.'.format(\n                ext_base_name=ext_base_name, flag=('was' if result else 'was NOT'),\n                fn_desc=fn_desc))\n\n        queue.put(result)\n\n    queue = Queue()\n    p = Process(target=_target_fn,\n                args=(ext_base_name, fn, fn_desc, queue, verbose))\n    p.daemon = True\n    p.start()\n    p.join()\n    return queue.get_nowait()\n\n\ndef extension_available(ext_base_name, verbose=False):\n    available_fn = lambda ext: ext is not None\n    return _check_extension_lambda(\n        ext_base_name, available_fn, 'built', verbose) or False\n\n\ndef mpi_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.mpi_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with MPI', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef gloo_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.gloo_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with Gloo', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef nccl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.nccl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with NCCL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef ddl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.ddl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with DDL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef mlsl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.mlsl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with MLSL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\n@contextmanager\ndef env(**kwargs):\n    # ignore args with None values\n    for k in list(kwargs.keys()):\n        if kwargs[k] is None:\n            del kwargs[k]\n\n    # backup environment\n    backup = {}\n    for k in kwargs.keys():\n        backup[k] = os.environ.get(k)\n\n    # set new values & yield\n    for k, v in kwargs.items():\n        os.environ[k] = v\n\n    try:\n        yield\n    finally:\n        # restore environment\n        for k in kwargs.keys():\n            if backup[k] is not None:\n                os.environ[k] = backup[k]\n            else:\n                del os.environ[k]\n", "path": "horovod/common/util.py"}, {"content": "# Copyright 2017 Uber Technologies, Inc. All Rights Reserved.\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\nimport horovod.tensorflow as hvd\nimport tensorflow as tf\n\n\ndef create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,\n                                 compression, sparse_as_dense):\n    class _DistributedOptimizer(keras.optimizers.Optimizer):\n        def __init__(self, name, device_dense, device_sparse, compression, sparse_as_dense,\n                     config):\n            if name is None:\n                name = \"Distributed%s\" % self.__class__.__base__.__name__\n            self._name = name\n            self._device_dense = device_dense\n            self._device_sparse = device_sparse\n            self._compression = compression\n            self._sparse_as_dense = sparse_as_dense\n            self._get_gradients_used = False\n            super(self.__class__, self).__init__(**config)\n\n        def get_gradients(self, loss, params):\n            \"\"\"\n            Compute gradients of all trainable variables.\n\n            See Optimizer.get_gradients() for more info.\n\n            In DistributedOptimizer, get_gradients() is overriden to also\n            allreduce the gradients before returning them.\n            \"\"\"\n            self._get_gradients_used = True\n            gradients = super(self.__class__, self).get_gradients(loss, params)\n            if hvd.size() > 1:\n                averaged_gradients = []\n                with tf.name_scope(self._name + \"_Allreduce\"):\n                    for grad in gradients:\n                        if grad is not None:\n                            if self._sparse_as_dense and \\\n                                    isinstance(grad, tf.IndexedSlices):\n                                grad = tf.convert_to_tensor(grad)\n                            avg_grad = hvd.allreduce(grad,\n                                                     device_dense=self._device_dense,\n                                                     device_sparse=self._device_sparse,\n                                                     compression=self._compression)\n                            averaged_gradients.append(avg_grad)\n                        else:\n                            averaged_gradients.append(None)\n                    return averaged_gradients\n            else:\n                return gradients\n\n        def apply_gradients(self, *args, **kwargs):\n            if not self._get_gradients_used:\n                  raise Exception('`apply_gradients()` was called without a call to '\n                                  '`get_gradients()`. If you\\'re using TensorFlow 2.0, '\n                                  'please specify `experimental_run_tf_function=False` in '\n                                  '`compile()`.')\n            return super(self.__class__, self).apply_gradients(*args, **kwargs)\n\n    # We dynamically create a new class that inherits from the optimizer that was passed in.\n    # The goal is to override get_gradients() method with an allreduce implementation.\n    # This class will have the same name as the optimizer it's wrapping, so that the saved\n    # model could be easily restored without Horovod.\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,),\n               dict(_DistributedOptimizer.__dict__))\n    return cls(name, device_dense, device_sparse, compression, sparse_as_dense,\n               optimizer.get_config())\n\n\ndef _eval(backend, op_or_result):\n    if hvd._executing_eagerly():\n        return op_or_result\n    else:\n        return backend.get_session().run(op_or_result)\n\n\nif hasattr(hvd, 'broadcast_global_variables'):\n    def broadcast_global_variables(backend, root_rank):\n        return _eval(backend, hvd.broadcast_global_variables(root_rank))\n\n\ndef allreduce(backend, value, name, average):\n    return _eval(backend, hvd.allreduce(tf.constant(value, name=name), average=average))\n\n\ndef allgather(backend, value, name):\n    return _eval(backend, hvd.allgather(tf.constant(value, name=name)))\n\n\ndef broadcast(backend, value, root_rank, name):\n    return _eval(backend, hvd.broadcast(tf.constant(value, name=name), root_rank))\n\n\ndef load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects):\n    horovod_objects = {\n        subclass.__name__.lower(): wrap_optimizer(subclass)\n        for subclass in keras.optimizers.Optimizer.__subclasses__()\n        if subclass.__module__ == keras.optimizers.Optimizer.__module__\n    }\n\n    if custom_optimizers is not None:\n        horovod_objects.update({\n            cls.__name__: wrap_optimizer(cls)\n            for cls in custom_optimizers\n        })\n\n    if custom_objects is not None:\n        horovod_objects.update(custom_objects)\n\n    return keras.models.load_model(filepath, custom_objects=horovod_objects)\n", "path": "horovod/_keras/__init__.py"}], "after_files": [{"content": "# Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n# Modifications copyright (C) 2019 Uber Technologies, Inc.\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\nfrom contextlib import contextmanager\nfrom multiprocessing import Process, Queue\nimport os\nimport sysconfig\n\nEXTENSIONS = ['tensorflow', 'torch', 'mxnet']\n\n\ndef get_ext_suffix():\n    \"\"\"Determine library extension for various versions of Python.\"\"\"\n    ext_suffix = sysconfig.get_config_var('EXT_SUFFIX')\n    if ext_suffix:\n        return ext_suffix\n\n    ext_suffix = sysconfig.get_config_var('SO')\n    if ext_suffix:\n        return ext_suffix\n\n    return '.so'\n\n\ndef get_extension_full_path(pkg_path, *args):\n    assert len(args) >= 1\n    dir_path = os.path.join(os.path.dirname(pkg_path), *args[:-1])\n    full_path = os.path.join(dir_path, args[-1] + get_ext_suffix())\n    return full_path\n\n\ndef check_extension(ext_name, ext_env_var, pkg_path, *args):\n    full_path = get_extension_full_path(pkg_path, *args)\n    if not os.path.exists(full_path):\n        raise ImportError(\n            'Extension %s has not been built.  If this is not expected, reinstall '\n            'Horovod with %s=1 to debug the build error.' % (ext_name, ext_env_var))\n\n\ndef _check_extension_lambda(ext_base_name, fn, fn_desc, verbose):\n    \"\"\"\n    Tries to load the extension in a new process.  If successful, puts fn(ext)\n    to the queue or False otherwise.  Mutes all stdout/stderr.\n    \"\"\"\n    def _target_fn(ext_base_name, fn, fn_desc, queue, verbose):\n        import importlib\n        import sys\n        import traceback\n\n        if verbose:\n            print('Checking whether extension {ext_base_name} was {fn_desc}.'.format(\n                ext_base_name=ext_base_name, fn_desc=fn_desc))\n        else:\n            # Suppress output\n            sys.stdout = open(os.devnull, 'w')\n            sys.stderr = open(os.devnull, 'w')\n\n        try:\n            ext = importlib.import_module('.' + ext_base_name, 'horovod')\n            result = fn(ext)\n        except:\n            traceback.print_exc()\n            result = None\n\n        if verbose:\n            print('Extension {ext_base_name} {flag} {fn_desc}.'.format(\n                ext_base_name=ext_base_name, flag=('was' if result else 'was NOT'),\n                fn_desc=fn_desc))\n\n        queue.put(result)\n\n    queue = Queue()\n    p = Process(target=_target_fn,\n                args=(ext_base_name, fn, fn_desc, queue, verbose))\n    p.daemon = True\n    p.start()\n    p.join()\n    return queue.get_nowait()\n\n\ndef extension_available(ext_base_name, verbose=False):\n    available_fn = lambda ext: ext is not None\n    return _check_extension_lambda(\n        ext_base_name, available_fn, 'built', verbose) or False\n\n\ndef mpi_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.mpi_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with MPI', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef gloo_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.gloo_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with Gloo', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef nccl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.nccl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with NCCL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef ddl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.ddl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with DDL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\ndef mlsl_built(verbose=False):\n    for ext_base_name in EXTENSIONS:\n        built_fn = lambda ext: ext.mlsl_built()\n        result = _check_extension_lambda(\n            ext_base_name, built_fn, 'built with MLSL', verbose)\n        if result is not None:\n            return result\n    return False\n\n\n@contextmanager\ndef env(**kwargs):\n    # ignore args with None values\n    for k in list(kwargs.keys()):\n        if kwargs[k] is None:\n            del kwargs[k]\n\n    # backup environment\n    backup = {}\n    for k in kwargs.keys():\n        backup[k] = os.environ.get(k)\n\n    # set new values & yield\n    for k, v in kwargs.items():\n        os.environ[k] = v\n\n    try:\n        yield\n    finally:\n        # restore environment\n        for k in kwargs.keys():\n            if backup[k] is not None:\n                os.environ[k] = backup[k]\n            else:\n                del os.environ[k]\n", "path": "horovod/common/util.py"}, {"content": "# Copyright 2017 Uber Technologies, Inc. All Rights Reserved.\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\nimport horovod.tensorflow as hvd\nimport tensorflow as tf\n\n\ndef create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,\n                                 compression, sparse_as_dense):\n    class _DistributedOptimizer(keras.optimizers.Optimizer):\n        def __init__(self, name, device_dense, device_sparse, compression, sparse_as_dense,\n                     config):\n            if name is None:\n                name = \"Distributed%s\" % self.__class__.__base__.__name__\n            self._name = name\n            self._device_dense = device_dense\n            self._device_sparse = device_sparse\n            self._compression = compression\n            self._sparse_as_dense = sparse_as_dense\n            self._get_gradients_used = False\n            super(self.__class__, self).__init__(**config)\n\n        def get_gradients(self, loss, params):\n            \"\"\"\n            Compute gradients of all trainable variables.\n\n            See Optimizer.get_gradients() for more info.\n\n            In DistributedOptimizer, get_gradients() is overriden to also\n            allreduce the gradients before returning them.\n            \"\"\"\n            self._get_gradients_used = True\n            gradients = super(self.__class__, self).get_gradients(loss, params)\n            if hvd.size() > 1:\n                averaged_gradients = []\n                with tf.name_scope(self._name + \"_Allreduce\"):\n                    for grad in gradients:\n                        if grad is not None:\n                            if self._sparse_as_dense and \\\n                                    isinstance(grad, tf.IndexedSlices):\n                                grad = tf.convert_to_tensor(grad)\n                            avg_grad = hvd.allreduce(grad,\n                                                     device_dense=self._device_dense,\n                                                     device_sparse=self._device_sparse,\n                                                     compression=self._compression)\n                            averaged_gradients.append(avg_grad)\n                        else:\n                            averaged_gradients.append(None)\n                    return averaged_gradients\n            else:\n                return gradients\n\n        def apply_gradients(self, *args, **kwargs):\n            if not self._get_gradients_used:\n                  raise Exception('`apply_gradients()` was called without a call to '\n                                  '`get_gradients()`. If you\\'re using TensorFlow 2.0, '\n                                  'please specify `experimental_run_tf_function=False` in '\n                                  '`compile()`.')\n            return super(self.__class__, self).apply_gradients(*args, **kwargs)\n\n        @classmethod\n        def from_config(cls, cfg):\n            return cls(name, device_dense, device_sparse, compression, sparse_as_dense, cfg)\n\n    # We dynamically create a new class that inherits from the optimizer that was passed in.\n    # The goal is to override get_gradients() method with an allreduce implementation.\n    # This class will have the same name as the optimizer it's wrapping, so that the saved\n    # model could be easily restored without Horovod.\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,),\n               dict(_DistributedOptimizer.__dict__))\n    return cls(name, device_dense, device_sparse, compression, sparse_as_dense,\n               optimizer.get_config())\n\n\ndef _eval(backend, op_or_result):\n    if hvd._executing_eagerly():\n        return op_or_result\n    else:\n        return backend.get_session().run(op_or_result)\n\n\nif hasattr(hvd, 'broadcast_global_variables'):\n    def broadcast_global_variables(backend, root_rank):\n        return _eval(backend, hvd.broadcast_global_variables(root_rank))\n\n\ndef allreduce(backend, value, name, average):\n    return _eval(backend, hvd.allreduce(tf.constant(value, name=name), average=average))\n\n\ndef allgather(backend, value, name):\n    return _eval(backend, hvd.allgather(tf.constant(value, name=name)))\n\n\ndef broadcast(backend, value, root_rank, name):\n    return _eval(backend, hvd.broadcast(tf.constant(value, name=name), root_rank))\n\n\ndef load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects):\n    horovod_objects = {\n        subclass.__name__.lower(): wrap_optimizer(subclass)\n        for subclass in keras.optimizers.Optimizer.__subclasses__()\n        if subclass.__module__ == keras.optimizers.Optimizer.__module__\n    }\n\n    if custom_optimizers is not None:\n        horovod_objects.update({\n            cls.__name__: wrap_optimizer(cls)\n            for cls in custom_optimizers\n        })\n\n    if custom_objects is not None:\n        horovod_objects.update(custom_objects)\n\n    return keras.models.load_model(filepath, custom_objects=horovod_objects)\n", "path": "horovod/_keras/__init__.py"}]}