Parallelizing across GPUs with jobs assigned by multiprocessing

[pvasired]

I have a multiprocessing.Pool which I would like to spin off several workers, each of which calls a function that uses JAX to perform operations on a specific GPU. My GPUs are not all the same model, so I can't use pmap(), I don't think?

Is this possible to do in JAX? Or is there an equivalent?

import os
import multiprocessing
import jax

def worker(gpu_id):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    return jax.devices()

if __name__ == "__main__":
    num_gpus = 2  # adjust this to match the number of available GPUs
    with multiprocessing.Pool(num_gpus) as p:
        p.map(worker, range(num_gpus))