code: def PartialDistributedOptimizer(
      optimizer,
      name=None,
      device_dense="",
      device_sparse="",
      compression=Compression.none,
      sparse_as_dense=False,
      gradient_predivide_factor=1.0,
      op=Average,
      backward_passes_per_step=1,
      average_aggregated_gradients=False,
      groups=None,
      process_set=global_process_set,
      local_layers=None,
      scale_local_gradients=True,
  ):
      """
      An optimizer that wraps another keras.optimizers.Optimizer, using an allreduce to
      average gradient values before applying gradients to model weights.
  
      Args:
          optimizer: Optimizer to use for computing gradients and applying updates.
          process_set: Gradients will only be reduced over Horovod processes belonging
                     to this process set. Defaults to the global process set.
          backward_passes_per_step: Number of backward passes to perform before calling
                                    hvd.allreduce. This allows accumulating updates over
                                    multiple mini-batches before reducing and applying them.
          average_aggregated_gradients: Whether to average the aggregated gradients that
                                        have been accumulated over multiple mini-batches.
                                        If true divides gradient updates by
                                        backward_passes_per_step.
                                        Only applicable for backward_passes_per_step > 1.
          local_layers: A collection of type tf.keras.layers.Layer local layers that their gradients need not
              to be synced accross ranks and is kept and applied locally.
              If not provided, the functionality of PartialDistributedOptimizer is
              identical to DistributedOptimizer.
          scale_local_gradients: Whether to scale the gradients of local variables. Default is set to True.
  
          The rest of the arguments are similar to those of DistributedOptimizer.
  
      """
      if gradient_predivide_factor != 1.0 and rocm_built():
          raise ValueError(
              "gradient_predivide_factor not supported yet with ROCm"
          )
  
      if op != Average and op != Sum:
          raise ValueError("op currently only supports Average and Sum")
  
      if groups is not None:
          if not (isinstance(groups, list) or groups > 0):
              raise ValueError(
                  "groups should be a non-negative integer or "
                  "a list of list of tf.Variable."
              )
  
      if local_layers is None:
          local_layers = []
      elif isinstance(local_layers, tf.keras.layers.Layer):
          local_layers = [local_layers]
      elif not all(
          isinstance(layer, tf.keras.layers.Layer)
          for layer in local_layers
      ):
          raise ValueError(
              "All local layers must be of tf.keras.layers.Layer type."
          )
  
      local_vars = [
          var
          for layer in local_layers
          for var in layer.trainable_weights
      ]
  
      _opt = _impl.create_distributed_optimizer(
          keras=keras,
          optimizer=optimizer,
          name=name,
          device_dense=device_dense,
          device_sparse=device_sparse,
          compression=compression,
          sparse_as_dense=sparse_as_dense,
          gradient_predivide_factor=gradient_predivide_factor,
          op=op,
          backward_passes_per_step=backward_passes_per_step,
          average_aggregated_gradients=average_aggregated_gradients,
          groups=groups,
          process_set=process_set,
          scale_local_gradients=scale_local_gradients,
      )
      if hasattr(_opt, "register_local_var"):
          for var in local_vars:
              _opt.register_local_var(var)
  
      return _opt
  
imports: import warnings
  from packaging import version
  import keras
  import keras.backend
  import tensorflow
  from horovod.tensorflow import init
  from horovod.tensorflow import shutdown
  from horovod.tensorflow import size, local_size, cross_size, rank, local_rank, cross_rank
  from horovod.tensorflow import is_initialized, start_timeline, stop_timeline
  from horovod.tensorflow import mpi_threads_supported, mpi_enabled, mpi_built
  from horovod.tensorflow import gloo_enabled, gloo_built
  from horovod.tensorflow import nccl_built, ddl_built, ccl_built, cuda_built, rocm_built
  from horovod.tensorflow import Average, Sum, Adasum
  from horovod.tensorflow.compression import Compression
  from horovod.tensorflow.mpi_ops import global_process_set
  from horovod.keras import callbacks, elastic
  import horovod._keras
label_desc: Use image node_id 2 for calling a global function with example usage: PartialDistributedOptimizer(optimizer, name, device_dense, device_sparse, compression, sparse_as_dense, gradient_predivide_factor, op, backward_passes_per_step, average_aggregated_gradients, groups, process_set, local_layers, scale_local_gradients) and returns: _opt