jax

[ROCm] Cherry-pick: Add ROCm xdist device pinning in conftest.py

#744

Merged