Feature Request: Native dynamic data loaders/generators

[mayureshagashe2105]

To my knowledge, there is no native support for this feature. Can the data loaders/ data generators be added which will dynamically read in the data at the training time using python's native generators?
These can have functionalities like:

• read from disk
• read from dataframe
• patch based learning
• etc.

Similar to keras.utils.Sequence or torch.utils.data.DataLoader

If yes, I would like to contribute.
Thank you.

[cgarciae]

Hey! This is not a official stance but I think that the JAX ecosystem has opted to use either tf.data, Pytorch DataLoaders, or any pipeline really as long as you can efficiently convert it into numpy and pass it to JAX. With tf.data is pretty easy using as_numpy_iterator, its also better tested as its used within Google. For Pytorch DataLoaders you can just convert to numpy first before feeding it into JAX.

[jlperla]

Curious if this unofficial stance has changed? I found https://github.com/birkhoffg/jax-dataloader to have a nice, simple interface (especially for pure JAX data). The alternatives of tf, pythorch, etc. are pretty massive dependencies and many introductory users of ML find the custom dataloader code a little scary.

Would a clean, simple JAX-only dataloader which uses the rngs=rngs arguments common to NNX be an acceptable small utility to add?