class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
code: def test_gradient_aggregation_with_local_vars(
      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
  
      total_num_variables = 8
      num_local_vars = 4
      X_0 = [0.0] * total_num_variables
      Y_0, Y_1 = [1.0] * total_num_variables, [
          2.0
      ] * total_num_variables
      X = [tf.Variable([x_0]) for x_0 in X_0]
      Y = [tf.Variable([y_0, y_1]) for y_0, y_1 in zip(Y_0, Y_1)]
      variables = [X, Y]
      flatten_variables = [
          item for sublist in variables for item in sublist
      ]
      for i in range(num_local_vars):
          x_var = X[i]
          y_var = Y[i]
          hvd_optimizer.register_local_var(x_var)
          hvd_optimizer.register_local_var(y_var)
  
      def loss():
          """
          loss = x - y * er where er = [float(hvd.rank()+1), 0.0]
          """
          # Gather the first row of y. It is equivalent to y * er.
          # Use tf.gather to produce tf.IndexedSlices gradient to improve test coverage.
          gathered_y_1 = tf.gather(Y, [0], axis=1)
          return X - (float(hvd.rank() + 1) * 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 = [-float(hvd.rank()+1), 0.0]
          Therefore, for each step, the value of x increases by 1.0 and
          the value of y increases by [-float(hvd.rank()), 0.0].
          """
          num_of_steps = (batch_id + 1) // backward_passes_per_step
  
          gradient_of_x = num_of_steps * 1.0
  
          # At each rank, the gradient of y_0 evaluates to (hvd.rank()+1).
          # For non-local variables we need to average the value accross the ranks.
          # The average value is (1+2+...hvd.size())/hvd.size() wich is
          # equivalent to the following expression:
          gradient_of_y_multiplier = float(hvd.size() + 1) / 2
          gradient_of_y_0 = (
              num_of_steps * gradient_of_y_multiplier * -1.0
          )
  
          if not average_aggregated_gradients:
              gradient_of_x *= backward_passes_per_step
              gradient_of_y_0 *= backward_passes_per_step
  
          expected_x = np.array(X_0)
  
          for i in range(len(X_0)):
              if i < num_local_vars:
                  expected_x[i] += gradient_of_x / hvd.size()
              else:
                  expected_x[i] += gradient_of_x
  
          expected_y_0 = np.array(Y_0)
          for i in range(len(Y_0)):
              if i < num_local_vars:
                  # recover the gradient of local vars
                  expected_y_0[i] += (
                      (gradient_of_y_0 / gradient_of_y_multiplier)
                      * float(hvd.rank() + 1)
                  ) / hvd.size()
              else:
                  expected_y_0[i] += gradient_of_y_0
  
          # It should remain constant as gradient is always 0.
          expected_y_1 = np.array(Y_1)
          expected_y = [
              [ey0, ey1] for ey0, ey1 in zip(expected_y_0, expected_y_1)
          ]
  
          return np.reshape(expected_x, (-1, 1)), np.array(expected_y)
  
      @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=flatten_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=flatten_variables,
                  experimental_aggregate_gradients=False,
              )
  
          expected_x, expected_y = compute_expected_value(idx)
          updated_x = np.array(variables[0])
          updated_y = np.array(variables[1])
          assert np.isclose(updated_x, expected_x).all()
          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 10 for calling the Tf2KerasTests obj's underlying member method code with example usage: obj.test_gradient_aggregation_with_local_vars(average_aggregated_gradients) and returns: unknown, np,  np, config, var