MPMD on multiple GPUs with JAX

[exenGT]

Hi,

I would like to write a program that performs the parallel tempering algorithm in JAX, but am not sure where to start. Specifically, the program needs to run two separate processes on two GPUs synchronously, and these two processes occasionally exchange data (say every 100 steps). If possible, could you provide some guidance on how this could be done?

Thank you,
JW

[MoFHeka]

any progress?