class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
code: def test_gradient_aggregation(self):
      with self.test_session(config=self.config) as sess:
  
          class TestingOptimizer(Optimizer):
              """
              Custom optimizer we use for testing gradient aggregation.
              """
  
              def get_config(self):
                  config = super().get_config()
                  return config
  
              def _create_slots(self, var_list):
                  pass
  
              def _resource_apply_dense(
                  self, grad, var, apply_state=None
              ):
                  return var.assign_add(grad)
  
          K.set_session(sess)
          session = tf.compat.v1.keras.backend.get_session(
              op_input_list=()
          )
  
          backward_passes_per_step = 4
          hvd_optimizer = hvd.DistributedOptimizer(
              optimizer=TestingOptimizer("test"),
              backward_passes_per_step=backward_passes_per_step,
              average_aggregated_gradients=True,
          )
          iterations = hvd_optimizer.iterations
          session.run(iterations.initializer)
  
          def compute_expected_value(batch_id):
              sum_per_aggregation = 0.0
              for _ in range(backward_passes_per_step):
                  grads_for_batch = 0.0
                  for rank in range(hvd.size()):
                      grads_for_batch += rank
  
                  # Apply `average_aggregated_gradients`.
                  grads_for_batch /= float(backward_passes_per_step)
  
                  # Averages across workers.
                  sum_per_aggregation += grads_for_batch / float(
                      hvd.size()
                  )
  
              aggregations_completed = math.floor(
                  (batch_id + 1) / backward_passes_per_step
              )
              return aggregations_completed * sum_per_aggregation
  
          grads = [tf.constant([float(hvd.rank())])]
          variables = [tf.Variable([0.0])]
          session.run(variables[0].initializer)
  
          allreduce_op = hvd_optimizer._allreduce(grads, variables)
          grads_and_vars = [(allreduce_op[0], variables[0])]
          apply_grads_op = hvd_optimizer.apply_gradients(grads_and_vars)
  
          for idx in range(10):
              _ = session.run(apply_grads_op)
  
              assert idx + 1 == session.run(hvd_optimizer.iterations)
              assert session.run(
                  variables[0].read_value()
              ) == compute_expected_value(idx)
  
imports: import math
  import numpy
  import os
  import pytest
  import sys
  import tensorflow
  import warnings
  from packaging import version
  from tensorflow import keras
  from horovod.common.util import is_version_greater_equal_than
  import horovod.tensorflow.keras
  from common import temppath
label_desc: Use image node_id 13 for calling the TfKerasTests obj's underlying member method code with example usage: obj.test_gradient_aggregation() and returns: unknown, config, var