Using pmap with one GPU

[ThomasWarford]

Hi all,
We have been using pmap to run the same function with different data on multiple CPU cores. However, when jax is run with GPU backend it runs the same programs with different data on each GPU.

We are wondering if there is a way to run the same program with different data on each GPU core. Is there a way to do this?

Any help is appreciated :)

Thanks

[ayaka14732]

@ThomasWarford There is no difference between the usage of pmap on CPU and GPU. So in your question, it run the same program with different data on each GPU core.