EpochIterator, TensorFlowTrainer : refactoring and fixes

keras/src/backend/jax/trainer.py

@@ -381,6 +381,7 @@ def fit(
         )
 
         self._symbolic_build(iterator=epoch_iterator)
+        epoch_iterator.reset()
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -405,7 +406,7 @@ def fit(
             callbacks.on_epoch_begin(epoch)
 
             self._jax_state_synced = True
-            for step, data in epoch_iterator.enumerate_epoch():
+            for step, data in epoch_iterator:
                 # Callbacks
                 callbacks.on_train_batch_begin(step)
 
@@ -539,6 +540,7 @@ def evaluate(
             )
 
         self._symbolic_build(iterator=epoch_iterator)
+        epoch_iterator.reset()
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -560,7 +562,7 @@ def evaluate(
         self.reset_metrics()
 
         self._jax_state_synced = True
-        for step, data in epoch_iterator.enumerate_epoch():
+        for step, data in epoch_iterator:
             callbacks.on_test_batch_begin(step)
 
             if self._jax_state_synced:
@@ -625,7 +627,7 @@ def predict(
 
         if not all(layer.built for layer in self._flatten_layers()):
             # Build the model on one batch of data.
-            for _, data in epoch_iterator.enumerate_epoch():
+            for _, data in epoch_iterator:
                 # Build model
                 x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data[0])
                 with backend.StatelessScope():
@@ -667,7 +669,7 @@ def append_to_outputs(batch_outputs, outputs):
         self._jax_state_synced = True
         outputs = None
         non_trainable_variables = None
-        for step, x in epoch_iterator.enumerate_epoch():
+        for step, x in epoch_iterator:
             callbacks.on_predict_batch_begin(step)
             if self._jax_state_synced:
                 # The state may have been synced by a callback.

keras/src/backend/numpy/trainer.py

@@ -212,7 +212,7 @@ def append_to_outputs(batch_outputs, outputs):
         self.stop_predicting = False
         callbacks.on_predict_begin()
         outputs = None
-        for step, data in epoch_iterator.enumerate_epoch():
+        for step, data in epoch_iterator:
             callbacks.on_predict_batch_begin(step)
             batch_outputs = self.predict_function(data)
             outputs = append_to_outputs(batch_outputs, outputs)
@@ -256,7 +256,7 @@ def evaluate(
 
         if not all(layer.built for layer in self._flatten_layers()):
             # Build the model on one batch of data.
-            for _, data in epoch_iterator.enumerate_epoch():
+            for _, data in epoch_iterator:
                 data_batch = data[0]
                 self._symbolic_build(data_batch)
                 break
@@ -278,7 +278,7 @@ def evaluate(
         callbacks.on_test_begin()
         logs = {}
         self.reset_metrics()
-        for step, data in epoch_iterator.enumerate_epoch():
+        for step, data in epoch_iterator:
             callbacks.on_test_batch_begin(step)
             logs = self.test_function(data)
             callbacks.on_test_batch_end(step, logs)

keras/src/backend/tensorflow/distribute_test.py

@@ -103,7 +103,7 @@ def test_epoch_iterator(self):
             distribute_strategy=strategy,
         )
         steps_seen = []
-        for step, data_iterator in epoch_iterator.enumerate_epoch():
+        for step, data_iterator in epoch_iterator:
             steps_seen.append(step)
             batch = next(data_iterator)
             self.assertEqual(len(batch), 3)