code: def create_distributed_optimizer(
      keras,
      optimizer,
      name,
      device_dense,
      device_sparse,
      compression,
      sparse_as_dense,
      gradient_predivide_factor,
      op,
      backward_passes_per_step=1,
      average_aggregated_gradients=False,
      groups=None,
      process_set=hvd.global_process_set,
      scale_local_gradients=True,
  ):
      class _DistributedOptimizer(*optimizer.__class__.__bases__):
          _HAS_AGGREGATE_GRAD = True
  
          def __init__(self, **kwargs):
              super(self.__class__, self).__init__(**kwargs)
  
              self._name = (
                  name
                  or "Distributed%s" % self.__class__.__base__.__name__
              )
              self._aggregated_gradients = False
  
              self._allreduce_grads = hvd._make_allreduce_grads_fn(
                  self._name,
                  device_dense,
                  device_sparse,
                  compression,
                  sparse_as_dense,
                  op,
                  gradient_predivide_factor,
                  groups,
                  process_set=process_set,
              )
  
              self._local_vars = set()
              self.process_set = process_set
              self.scale_local_gradients = scale_local_gradients
              self._agg_helper = None
              if backward_passes_per_step > 1:
                  if hvd._executing_eagerly():
                      self._agg_helper = LocalGradientAggregationHelperEager(
                          backward_passes_per_step=backward_passes_per_step,
                          allreduce_func=self._allreduce_grads,
                          sparse_as_dense=sparse_as_dense,
                          average_aggregated_gradients=average_aggregated_gradients,
                          process_set=process_set,
                          scale_local_gradients=scale_local_gradients,
                      )
                  else:
                      self._agg_helper = LocalGradientAggregationHelper(
                          backward_passes_per_step=backward_passes_per_step,
                          allreduce_func=self._allreduce_grads,
                          sparse_as_dense=sparse_as_dense,
                          average_aggregated_gradients=average_aggregated_gradients,
                          rank=rank(),
                          optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_KERAS,
                          process_set=process_set,
                          scale_local_gradients=scale_local_gradients,
                      )
  
          if version.parse(tf.__version__) >= version.parse("2.12.0"):
  
              @property
              def variables(self):
                  return CallableList(
                      super(self.__class__, self).variables()
                  )
  
          else:
  
              def variables(self):
                  if _IS_TF2:
                      return super(self.__class__, self).variables()
                  return self.get_weights()
  
          def register_local_var(self, var):
              """Registers a source/variable as worker local. Horovod will not perform any global
              operations on gradients corresponding to these sources and will instead return the local
              gradient."""
              if self._agg_helper:
                  self._agg_helper.register_local_var(var)
              elif _IS_TF2:
                  self._local_vars.add(var.ref())
              else:
                  self._local_vars.add(var)
  
          def compute_gradients(self, loss, var_list, tape=None):
              return self._compute_gradients(loss, var_list, None, tape)
  
          def _compute_gradients(
              self, loss, var_list, grad_loss=None, tape=None
          ):
              """
              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.
              """
              base_class = super(self.__class__, self)
              if _PRE_TF_2_4_0:
                  return base_class._compute_gradients(
                      loss, var_list, grad_loss, tape
                  )
  
              tape = tf.GradientTape() if tape is None else tape
              if hasattr(base_class, "_compute_gradients"):
                  grads_and_vars = base_class._compute_gradients(
                      # pylint: disable=protected-access
                      loss,
                      var_list,
                      grad_loss,
                      tape=tape,
                  )
              else:
                  grads_and_vars = base_class.compute_gradients(
                      # pylint: disable=protected-access
                      loss,
                      var_list,
                      tape=tape,
                  )
              grads, weights = list(zip(*grads_and_vars))
  
              allreduced_grads = self._allreduce(grads, weights)
              return list(zip(allreduced_grads, weights))
  
          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.
              """
              gradients = super(self.__class__, self).get_gradients(
                  loss, params
              )
              return self._allreduce(gradients, params)
  
          def _aggregate_gradients(self, grads_and_vars):
              base_class = super(self.__class__, self)
              if _PRE_TF_2_4_0:
                  grads, vars = list(zip(*grads_and_vars))
                  aggregated_grads = self._allreduce(grads, vars)
                  return aggregated_grads
              elif hasattr(base_class, "_aggregate_gradients"):
                  return base_class._aggregate_gradients(grads_and_vars)
              else:
                  return base_class.aggregate_gradients(grads_and_vars)
  
          def _allreduce(self, grads, vars):
              self._aggregated_gradients = True
  
              if self._agg_helper:
                  return self._agg_helper.compute_gradients(
                      tuple(grads), tuple(vars)
                  )
              else:
  
                  def __filtered_reduce_grads(grads, vars):
                      rv = []
                      rg = []
                      if _IS_TF2:
                          v2g = {
                              var.ref(): grad
                              for var, grad in zip(vars, grads)
                          }
                          for var, grad in zip(vars, grads):
                              if var.ref() not in self._local_vars:
                                  rv.append(var)
                                  rg.append(grad)
                      else:
                          v2g = {
                              var: grad
                              for var, grad in zip(vars, grads)
                          }
                          for var, grad in zip(vars, grads):
                              if var not in self._local_vars:
                                  rv.append(var)
                                  rg.append(grad)
  
                      rg = self._allreduce_grads(rg, rv)
                      horovod_size = (
                          size_op(
                              process_set_id=self.process_set.process_set_id
                          )
                          if int(os.environ.get("HOROVOD_ELASTIC", 0))
                          else self.process_set.size()
                      )
                      if _IS_TF2:
                          for rv, rg in zip(rv, rg):
                              v2g[rv.ref()] = rg
  
                          if self.scale_local_gradients and len(
                              self._local_vars
                          ):
                              # Scale local gradients by a size factor. See pull/3695 and discussions/3705 for context.
                              for v_ref in v2g:
                                  if (
                                      v_ref in self._local_vars
                                      and v2g[v_ref] is not None
                                  ):
                                      grad = v2g[v_ref]
                                      if isinstance(
                                          grad, tf.IndexedSlices
                                      ):
                                          grad = tf.IndexedSlices(
                                              grad.values
                                              / float(horovod_size),
                                              grad.indices,
                                              grad.dense_shape,
                                          )
                                      else:
                                          grad /= float(horovod_size)
                                      v2g[v_ref] = grad
  
                          return [v2g[rv.ref()] for rv in vars]
                      else:
                          for rv, rg in zip(rv, rg):
                              v2g[rv] = rg
  
                          if self.scale_local_gradients and len(
                              self._local_vars
                          ):
                              # Scale local gradients by a size factor. See pull/3695 and discussions/3705 for context.
                              for v in v2g:
                                  if (
                                      v in self._local_vars
                                      and v2g[v] is not None
                                  ):
                                      grad = v2g[v]
                                      if isinstance(
                                          grad, tf.IndexedSlices
                                      ):
                                          grad = tf.IndexedSlices(
                                              grad.values
                                              / float(horovod_size),
                                              grad.indices,
                                              grad.dense_shape,
                                          )
                                      else:
                                          grad /= float(horovod_size)
                                      v2g[v] = grad
  
                          return [v2g[rv] for rv in vars]
  
                  return __filtered_reduce_grads(grads, vars)
  
          def apply_gradients(self, *args, **kwargs):
              if self._agg_helper:
                  if isinstance(args[0], zip):
                      # If grad_and_vars are passed in as a zip object
                      # convert to a list. This is necessary for TF2.4+
                      # b/c args[0] is used in both conditional branches
                      # inside _agg_helper.apply_gradients().
                      args = list(args)
                      args[0] = list(args[0])
                      args = tuple(args)
  
                  results = self._agg_helper.apply_gradients(
                      lambda: super(
                          self.__class__, self
                      ).apply_gradients(*args, **kwargs),
                      self,
                      *args,
                      **kwargs,
                  )
              else:
                  results = super(self.__class__, self).apply_gradients(
                      *args, **kwargs
                  )
  
              if _PRE_TF_2_4_0 and not self._aggregated_gradients:
                  raise Exception(
                      "`apply_gradients()` was called without a call to "
                      "`get_gradients()` or `_aggregate_gradients`. If you're "
                      "using TensorFlow 2.0, please specify "
                      "`experimental_run_tf_function=False` in `compile()`."
                  )
  
              return results
  
      # 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())
  
file_path: horovod/horovod/_keras/__init__.py
imports: import os
  from packaging import version
  import horovod.tensorflow
  import tensorflow
  from horovod.tensorflow.gradient_aggregation import LocalGradientAggregationHelper
  from horovod.tensorflow.gradient_aggregation_eager import LocalGradientAggregationHelperEager
  from horovod.tensorflow.mpi_ops import rank, size_op
label_desc: Use image node_id 1 for calling a global function with example usage: create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse, compression, sparse_as_dense, gradient_predivide_factor, op, backward_passes_per_step, average_aggregated_gradients, groups, process_set, scale_local_gradients) and returns: cls, self, list, self, results, CallableList, self, base_class, aggregated_grads, self, __filtered_reduce_grads, super, base_class, base_class, unknown, unknown