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, average_aggregated_gradients):
      class TestingOptimizer(Optimizer):
          """
          Custom optimizer we use for testing gradient aggregation.
          """
  
          def __init__(self, name, **kwargs):
              super().__init__(name=name, **kwargs)
              if hasattr(self, "_build_learning_rate"):
                  self._learning_rate = self._build_learning_rate(0.1)
  
          def get_config(self):
              config = super().get_config()
              return config
  
          def _create_slots(self, var_list):
              # Only needed for TF < 2.2.
              pass
  
          def _resource_apply_dense(self, grad, var, apply_state=None):
              return var.assign_add(grad)
  
          def update_step(self, gradient, variable):
              variable.assign_add(gradient)
  
      backward_passes_per_step = 4
      hvd_optimizer = hvd.DistributedOptimizer(
          optimizer=TestingOptimizer("test"),
          backward_passes_per_step=backward_passes_per_step,
          average_aggregated_gradients=average_aggregated_gradients,
          sparse_as_dense=True,
      )
      _ = hvd_optimizer.iterations
  
      x_0 = 0.0
      y_0, y_1 = 1.0, 2.0
      x = tf.Variable([x_0])
      y = tf.Variable([y_0, y_1])
      variables = [x, y]
  
      def loss():
          """
          loss = x - y * e1 where e1 = [1.0, 0.0]
          """
          # Gather the first row of y. It is equivalent to y * e1.
          # Use tf.gather to produce tf.IndexedSlices gradient to improve test coverage.
          gathered_y_1 = tf.gather(y, [0])
          return x - gathered_y_1
  
      def compute_expected_value(batch_id):
          """
          Given the loss function above, the gradient of x and y can be derived.
            dloss/dx = 1.0
            dloss/dy = [-1.0, 0.0]
          Therefore, for each step, the value of x increases by 1 and
          the value of y increases by [-1.0, 0.0].
          """
          num_of_steps = (batch_id + 1) // backward_passes_per_step
  
          gradient_of_x = num_of_steps * 1.0
          gradient_of_y_0 = num_of_steps * -1.0
  
          if not average_aggregated_gradients:
              gradient_of_x *= backward_passes_per_step
              gradient_of_y_0 *= backward_passes_per_step
  
          # Add gradient with its initial values because the TestingOptimizer optimize
          # variable by assign_add.
          expected_x = x_0 + gradient_of_x
          expected_y_0 = y_0 + gradient_of_y_0
          # It should remain constant as gradient is always 0.
          expected_y_1 = y_1
  
          return np.array(expected_x), np.array(
              [expected_y_0, expected_y_1]
          )
  
      @tf.function
      def compute_and_apply_gradients_in_tf_function(
          var_list, **kwargs
      ):
          # Compute and apply gradient updates in tf.function to reproduce
          # how it is done inside `model.fit()`.
          grads_and_vars = hvd_optimizer._compute_gradients(
              loss, var_list=var_list
          )
          hvd_optimizer.apply_gradients(grads_and_vars, **kwargs)
  
      total_num_of_steps = 10
      for idx in range(total_num_of_steps):
          if _PRE_TF_2_2_0:
              compute_and_apply_gradients_in_tf_function(
                  var_list=variables
              )
          else:
              # In 2.2 and 2.3 the horovod optimizer sets `_HAS_AGGREGATE_GRAD = True`.
              # This configures tf.keras to call `_aggregate_gradients()` outside of
              # `apply_gradients()` and to set `experimental_aggregate_gradients` to
              # False when calling `apply_gradients()` to prevent it from calling
              # `_aggregate_gradients()` again.
              compute_and_apply_gradients_in_tf_function(
                  var_list=variables,
                  experimental_aggregate_gradients=False,
              )
  
          expected_x, expected_y = compute_expected_value(idx)
          updated_x = variables[0].numpy()
          updated_y = variables[1].numpy()
          assert np.isclose(updated_x, expected_x)
          assert np.isclose(updated_y, expected_y).all()
          assert idx + 1 == hvd_optimizer.iterations.numpy()
  
      aggregation_counter = hvd_optimizer._agg_helper.counter.numpy()
      assert (
          aggregation_counter
          == total_num_of_steps % backward_passes_per_step
      )
  
imports: import tensorflow
  import numpy
  import warnings
  from packaging import version
  from parameterized import parameterized
  import pytest
  from tensorflow import keras
  from horovod.common.util import is_version_greater_equal_than
  import horovod.keras
  import horovod.tensorflow.keras
label_desc: Use image node_id 9 for calling the Tf2KerasTests obj's underlying member method code with example usage: obj.test_gradient_aggregation(average_aggregated_gradients) and returns: unknown, np,  np, config, var