class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
docstr: Note: test makes most sense with more than 1 nodes. 
docstr_tok: ["Note",":","test","makes","most","sense","with","more","than","1","nodes","."]
code: def test_partial_distributed_optimizer(self):
      """Note: test makes most sense with more than 1 nodes."""
      if hvd.size() == 1:
          self.skipTest("Only one worker available")
  
      # the keras model has 3 layers, we test cases with 0, 1, and 2 local layers.
      for num_local_layers in range(3):
          model = tf.keras.models.Sequential()
          initializer = tf.keras.initializers.Constant(hvd.rank())
          model.add(
              tf.keras.layers.Dense(
                  2,
                  input_shape=(3,),
                  kernel_initializer=initializer,
                  bias_initializer=initializer,
              )
          )
          model.add(tf.keras.layers.RepeatVector(3))
          model.add(
              tf.keras.layers.TimeDistributed(
                  tf.keras.layers.Dense(
                      3,
                      kernel_initializer=initializer,
                      bias_initializer=initializer,
                  )
              )
          )
          opt = tf.keras.optimizers.Adam()
          model.compile(
              loss=tf.keras.losses.MSE,
              metrics=[tf.keras.metrics.categorical_accuracy],
          )
  
          X = np.random.random((1, 3))
          Y = np.random.random((1, 3, 3))
  
          try:
              init = tf.global_variables_initializer()
          except AttributeError:
              init = tf.compat.v1.global_variables_initializer()
          self.evaluate(init)
  
          with tf.GradientTape(persistent=True) as tape:
              p = model(X, training=True)
              l = model.loss(Y, p)
  
          gradients_tape = tape.gradient(l, model.trainable_weights)
  
          # deem local layers
          local_layers = model.layers[:num_local_layers]
  
          opt = hvd_keras.PartialDistributedOptimizer(
              opt, sparse_as_dense=True, local_layers=local_layers
          )
          gradients_vars_opt = opt._compute_gradients(
              l, model.trainable_weights, tape=tape
          )
  
          var_grad_tape = {
              var.ref(): grad
              for var, grad in zip(
                  model.trainable_weights, gradients_tape
              )
          }
          var_grad_opt = {
              var.ref(): grad for grad, var in gradients_vars_opt
          }
          local_vars = [
              var.ref()
              for layer in local_layers
              for var in layer.trainable_weights
          ]
  
          for var in model.trainable_weights:
              if var.ref() in local_vars:
                  # local gradients should not change.
                  self.assertAllClose(
                      var_grad_tape[var.ref()],
                      hvd.size() * var_grad_opt[var.ref()],
                  )
              else:
                  # non-local gradients shouldn't be equal given that the initial weights are set to ranks
                  self.assertNotAllClose(
                      var_grad_tape[var.ref()], var_grad_opt[var.ref()]
                  )
  
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 11 for calling the Tf2KerasTests obj's underlying member method code with example usage: obj.test_partial_distributed_optimizer() without return types