Running tests in parallel in a multi-GPU machine (one test per GPU)

[reza-amd]

Hi all,
I wanted to know you recommendation on running unit tests in parallel (for a multi-gpu machine), but only run a single test per GPU at a time.
In TensorFlow, you use bazel "run_under" argument with this script to control the number of unit tests running in parallel per GPU.
I appreciate if you can provide more information about your solution for JAX in your internal CI builds to provide the same controllability with "pytest".

[reza-amd]

This question has also been asked in issue #7323

[hawkinsp]

I replied in #7323 !