Using multiple streams on a single GPU

[alonkukl]

I'm using JAX to speed up some linear algebra calculations by JITing them and running them on a GPU.

My issue is that I have a big time overhead preparing the data for the GPU, thus it is idle long periods of time.

I was wondering if I could use different GPU streams on different JAX processes by that "hiding" the overhead. Meaning using only a single GPU in pipeline between the different processes.

I went over the documentation but could find a way to do it.