# Buggy code
def create_distributed_optimizer(
    keras, optimizer, name, device_dense, device_sparse, compression, sparse_as_dense
):
    class _DistributedOptimizer(keras.optimizers.Optimizer):
        def __init__(
            self,
            name,
            device_dense,
            device_sparse,
            compression,
            sparse_as_dense,
            config,
        ):
            if name is None:
                name = "Distributed%s" % self.__class__.__base__.__name__
            self._name = name
            self._device_dense = device_dense
            self._device_sparse = device_sparse
            self._compression = compression
            self._sparse_as_dense = sparse_as_dense
            self._get_gradients_used = False
            super(self.__class__, self).__init__(**config)

        def get_gradients(self, loss, params):
            """
            Compute gradients of all trainable variables.

            See Optimizer.get_gradients() for more info.

            In DistributedOptimizer, get_gradients() is overriden to also
            allreduce the gradients before returning them.
            """
            self._get_gradients_used = True
            gradients = super(self.__class__, self).get_gradients(loss, params)
            if hvd.size() > 1:
                averaged_gradients = []
                with tf.name_scope(self._name + "_Allreduce"):
                    for grad in gradients:
                        if grad is not None:
                            if self._sparse_as_dense and isinstance(
                                grad, tf.IndexedSlices
                            ):
                                grad = tf.convert_to_tensor(grad)
                            avg_grad = hvd.allreduce(
                                grad,
                                device_dense=self._device_dense,
                                device_sparse=self._device_sparse,
                                compression=self._compression,
                            )
                            averaged_gradients.append(avg_grad)
                        else:
                            averaged_gradients.append(None)
                    return averaged_gradients
            else:
                return gradients

        def apply_gradients(self, *args, **kwargs):
            if not self._get_gradients_used:
                raise Exception(
                    "`apply_gradients()` was called without a call to "
                    "`get_gradients()`. If you're using TensorFlow 2.0, "
                    "please specify `experimental_run_tf_function=False` in "
                    "`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
    # model could be easily restored without Horovod.
    cls = type(
        optimizer.__class__.__name__,
        (optimizer.__class__,),
        dict(_DistributedOptimizer.__dict__),
    )
    return cls(
        name,
        device_dense,
        device_sparse,
        compression,
        sparse_as_dense,
        optimizer.get_config(),
    )


# Traceback (AttributeError)
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 541, in __getattribute__
[1,0]<stderr>:    return super(OptimizerV2, self).__getattribute__(name)
[1,0]<stderr>:AttributeError: 'Adam' object has no attribute 'lr'
[1,0]<stderr>:
[1,0]<stderr>:During handling of the above exception, another exception occurred:
[1,0]<stderr>:
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "tensorflow_keras_mnist.py", line 94, in <module>
[1,0]<stderr>:    mnist_model.fit(dataset, steps_per_epoch=500 // hvd.size(), callbacks=callbacks, epochs=24, verbose=verbose)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
[1,0]<stderr>:    use_multiprocessing=use_multiprocessing)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py", line 693, in fit
[1,0]<stderr>:    steps_name='steps_per_epoch')
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py", line 196, in model_iteration
[1,0]<stderr>:    callbacks._call_begin_hook(mode)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/callbacks.py", line 249, in _call_begin_hook
[1,0]<stderr>:    self.on_train_begin()
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/callbacks.py", line 365, in on_train_begin
[1,0]<stderr>:    callback.on_train_begin(logs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/horovod/_keras/callbacks.py", line 137, in on_train_begin
[1,0]<stderr>:    self.initial_lr = self.backend.get_value(self.model.optimizer.lr)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 550, in __getattribute__
[1,0]<stderr>:    return self._get_hyper(name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 527, in _get_hyper
[1,0]<stderr>:    self._create_hypers()
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 643, in _create_hypers
[1,0]<stderr>:    aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 805, in add_weight
[1,0]<stderr>:    aggregation=aggregation)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 744, in _add_variable_with_custom_getter
[1,0]<stderr>:    **kwargs_for_getter)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer_utils.py", line 139, in make_variable
[1,0]<stderr>:    shape=variable_shape if variable_shape else None)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 258, in __call__
[1,0]<stderr>:    return cls._variable_v1_call(*args, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 219, in _variable_v1_call
[1,0]<stderr>:    shape=shape)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 197, in <lambda>
[1,0]<stderr>:    previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 2507, in default_variable_creator
[1,0]<stderr>:    shape=shape)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 262, in __call__
[1,0]<stderr>:    return super(VariableMetaclass, cls).__call__(*args, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py", line 1406, in __init__
[1,0]<stderr>:    distribute_strategy=distribute_strategy)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/resourc[1,0]<stderr>:e_variable_ops.py", line 1538, in _init_from_args
[1,0]<stderr>:    name="initial_value", dtype=dtype)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1184, in convert_to_tensor
[1,0]<stderr>:    return convert_to_tensor_v2(value, dtype, preferred_dtype, name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1242, in convert_to_tensor_v2
[1,0]<stderr>:    as_ref=False)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1296, in internal_convert_to_tensor
[1,0]<stderr>:    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 286, in _constant_tensor_conversion_function
[1,0]<stderr>:    return constant(v, dtype=dtype, name=name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 227, in constant
[1,0]<stderr>:    allow_broadcast=True)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 235, in _constant_impl
[1,0]<stderr>:    t = convert_to_eager_tensor(value, ctx, dtype)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 96, in convert_to_eager_tensor
[1,0]<stderr>:    return ops.EagerTensor(value, ctx.device_name, dtype)
[1,0]<stderr>:ValueError: Attempt to convert a value ({'class_name': 'ExponentialDecay', 'config': {'initial_learning_rate': 0.001, 'decay_steps': 100000, 'decay_rate': 0.96, 'staircase': True, 'name': None}}) with an unsupported type (<class 'dict'>) to a Tensor.
[1,0]<stderr>:Exception ignored in: <function _RandomSeedGeneratorDeleter.__del__ at 0x7f7dec944a70>
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3009, in __del__
[1,0]<stderr>:AttributeError: 'NoneType' object has no attribute 'device'
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------

# Fixed code
def create_distributed_optimizer(
    keras, optimizer, name, device_dense, device_sparse, compression, sparse_as_dense
):
    class _DistributedOptimizer(keras.optimizers.Optimizer):
        def __init__(self, **kwargs):
            self._name = name or "Distributed%s" % self.__class__.__base__.__name__
            self._device_dense = device_dense
            self._device_sparse = device_sparse
            self._compression = compression
            self._sparse_as_dense = sparse_as_dense
            self._get_gradients_used = False
            super(self.__class__, self).__init__(**kwargs)

        def get_gradients(self, loss, params):
            """
            Compute gradients of all trainable variables.

            See Optimizer.get_gradients() for more info.

            In DistributedOptimizer, get_gradients() is overriden to also
            allreduce the gradients before returning them.
            """
            self._get_gradients_used = True
            gradients = super(self.__class__, self).get_gradients(loss, params)
            if hvd.size() > 1:
                averaged_gradients = []
                with tf.name_scope(self._name + "_Allreduce"):
                    for grad in gradients:
                        if grad is not None:
                            if self._sparse_as_dense and isinstance(
                                grad, tf.IndexedSlices
                            ):
                                grad = tf.convert_to_tensor(grad)
                            avg_grad = hvd.allreduce(
                                grad,
                                device_dense=self._device_dense,
                                device_sparse=self._device_sparse,
                                compression=self._compression,
                            )
                            averaged_gradients.append(avg_grad)
                        else:
                            averaged_gradients.append(None)
                    return averaged_gradients
            else:
                return gradients

        def apply_gradients(self, *args, **kwargs):
            if not self._get_gradients_used:
                raise Exception(
                    "`apply_gradients()` was called without a call to "
                    "`get_gradients()`. If you're using TensorFlow 2.0, "
                    "please specify `experimental_run_tf_function=False` in "
                    "`compile()`."
                )
            return super(self.__class__, self).apply_gradients(*args, **kwargs)

    # 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
    # model could be easily restored without Horovod.
    cls = type(
        optimizer.__class__.__name__,
        (optimizer.__class__,),
        dict(_DistributedOptimizer.__dict__),
    )
    return cls.from_config(optimizer.get_config())