Most efficient way to get data in JAX

[thomashirtz]

Hello!
I would like to develop program with very efficient use of GPU. I learned how to handle the training part of neural network. Now I want to know if there is a way to optimize as much as possible the data retrieval part of the function. Is there a solution to be very efficient for retrieving the data needed to train the neural networks ? (currently using ray.Queue with processes putting a lot of data in it to be sure that it is always furnished) I imagine that this part cannot be jitted, even when using shared_memory spaces or else ?
Also is there a very efficient way to send the weight of the current network to other processes ?
Thanks