Does Flax support DataGenerators?

[halixness]

Hello,
I am currently developing a video prediction model with Keras under Tensorflow 2.3.1. I have developed a Keras DataGenerator (inheriting keras.utils.Sequence) to generate batches by reading the necessary data files "on the fly", from a very large selection.
I would like to migrate to Flax and I was wondering if keras.Sequence generators are supported.
Sorry for the simple question, thank you very much!

[avital]

Hi @halixness! JAX and Flax are completely agnostic to data loading pipelines, folks have successfully used Flax with Torch dataloaders, for example. So if you take any of our examples and change the data loader code it should continue to work.