`EpochIterator` initialized with a generator cannot be reset

[nicolaspi]

EpochIterator initialized with a generator cannot be reset, consequently the following bug might occur: #20394
Another consequence is running prediction on an unbuilt model discards samples:

import os
import numpy as np
os.environ["KERAS_BACKEND"] = "jax"
import keras

set_size = 4
def generator():
  for _ in range(set_size):
    yield np.ones((1, 4)), np.ones((1, 1))

model = keras.models.Sequential([keras.layers.Dense(1, activation="relu")])
model.compile(loss="mse", optimizer="adam")
y = model.predict(generator())

print(f"There should be {set_size} predictions, there are {len(y)}.")

There should be 4 predictions, there are 2.

[mehtamansi29]

Hi @nicolaspi -

Thanks for reporting the issue. Here generator() function yield values once at a time. And in jax backend, model.predict() is compiled using jit compilation. So generator() will not able to provide data to model.predict() via jit compilation process.

Here if you will train the model using fit method then model.predict() will give 4 predictions with jax backend.
Attached gist for the reference.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No