Getting rank of device within pmapped function

[minqi]

Is there anything equivalent to torch.distributed.get_rank() in JAX?

My use case is that I would like to selectively access a different index of a nested pytree inside a static argument of my pmapped function, depending on the corresponding device index.

[minqi]

Answering my own question: Just realized that's what jax.lax.axis_index is for!