jax icon indicating copy to clipboard operation
jax copied to clipboard

[ROCm] Disable CUDA specific test for ROCm.

Open rsanthanam-amd opened this issue 3 years ago • 1 comments

/cc @hawkinsp

rsanthanam-amd avatar Aug 27 '22 12:08 rsanthanam-amd

@hawkinsp gentle ping

rsanthanam-amd avatar Sep 08 '22 12:09 rsanthanam-amd

Sorry this dropped off my list of things to do. Apologies!

I had missed that multiprocess support isn't hooked into ROCM/RCCL (which is what nccl_unique_id_callback does). I'd guess that's a fairly mechanical thing to fix but since that's the state of things I agree with your original fix: disable the test.

https://github.com/google/jax/pull/12556 fixes.

hawkinsp avatar Sep 28 '22 13:09 hawkinsp