code: def test_epoch_iterator(self):
      x = np.random.random((100, 16))
      y = np.random.random((100, 4))
      sample_weight = np.random.random((100,))
      batch_size = 16
      shuffle = True
  
      strategy = tf.distribute.MirroredStrategy(["CPU:0", "CPU:1"])
  
      epoch_iterator = tf_trainer.TFEpochIterator(
          x=x,
          y=y,
          sample_weight=sample_weight,
          batch_size=batch_size,
          shuffle=shuffle,
          distribute_strategy=strategy,
      )
      steps_seen = []
      for step, data_iterator in epoch_iterator.enumerate_epoch():
          steps_seen.append(step)
          batch = next(data_iterator)
          self.assertEqual(len(batch), 3)
          x, y, sample_weight = batch
          self.assertTrue(
              isinstance(
                  x, tf.types.experimental.distributed.PerReplica
              )
          )
          # Make sure the local batch size is 8
          if step < 6:
              self.assertEqual(x.values[0].shape, [8, 16])
              self.assertEqual(y.values[0].shape, [8, 4])
              self.assertEqual(sample_weight.values[0].shape, [8])
          else:
              # Last partial batch
              self.assertEqual(x.values[0].shape, [2, 16])
              self.assertEqual(y.values[0].shape, [2, 4])
              self.assertEqual(sample_weight.values[0].shape, [2])
      self.assertEqual(steps_seen, [0, 1, 2, 3, 4, 5, 6])
  
file_path: keras/keras/backend/tensorflow/distribute_test.py
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow.python.eager import context
  from keras import backend
  from keras import layers
  from keras import models
  from keras import testing
  from keras.backend.tensorflow import trainer
label_desc: Use image node_id 4 for calling the DistributeTest obj's underlying member method code with example usage: obj.test_epoch_iterator() without return types