class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
code: def test_elastic_state(self):
      with self.test_session(config=self.config) as sess:
          K.set_session(sess)
  
          v = 1.0 if hvd.rank() == 0 else 2.0
          model1 = tf.keras.Sequential(
              [tf.keras.layers.Dense(2, activation="softmax")]
          )
          model1.build((2, 2))
          model1.set_weights(
              [
                  np.array([[v, v], [v, v]], dtype=np.float32),
                  np.array([v, v], dtype=np.float32),
              ]
          )
  
          model2 = tf.keras.Sequential(
              [tf.keras.layers.Dense(2, activation="softmax")]
          )
          model2.build((2, 2))
          model2.set_weights(
              [
                  np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
                  np.array([0.0, 0.0], dtype=np.float32),
              ]
          )
  
          optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())
  
          state = hvd.elastic.KerasState(
              model1,
              optimizer,
              batch=20 + hvd.rank(),
              epoch=10 + hvd.rank(),
          )
          state.sync()
  
          model1_weights = model1.get_weights()
          model2_weights = model2.get_weights()
  
          # After sync, all values should match the root rank
          for w in state.model.get_weights():
              self.assertAllClose(w, np.ones_like(w))
          assert state.batch == 20
          assert state.epoch == 10
  
          # Partially modify then restore
          model1.set_weights(model2_weights)
          state.batch = 21
          state.epoch = 11
  
          state.restore()
  
          for w1, w2 in zip(model1.get_weights(), model1_weights):
              self.assertAllClose(w1, w2)
          assert state.batch == 20
          assert state.epoch == 10
  
          # Partially modify then commit
          model1.set_weights(model2_weights)
          state.batch = 21
          state.epoch = 11
  
          state.commit()
          state.restore()
  
          for w1, w2 in zip(model1.get_weights(), model2_weights):
              self.assertAllClose(w1, w2)
          assert state.batch == 21
          assert state.epoch == 11
  
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 12 for calling the TfKerasTests obj's underlying member method code with example usage: obj.test_elastic_state() without return types