How to get the device id in a multi-gpu host

[jrabary]

Dear all,

I need to get the id of the device where my computation takes place.

Say for example, I have 8 gpu on my host and I need to multiply a tensor by a specific number per GPU
like a * (GPU_ID + 1).

There is a jax.device_count but I'd rather need a jax.device_id or jax.local_device_id

[jrabary]

More precisely, I'm looking for the jax implementation of the following tensorflow function

replica_context = tf.distribute.get_replica_context()
id = replica_context.replica_id_in_sync_group

[skye]

I assume you're using pmap? One option is to pass jnp.arange(jax.device_count()) into the pmap, so every GPU will get a different value. Alternatively, you can use jax.lax.axis_index.

[jrabary]

Thank you, this is what I need.