jaxlib error on blackwell GPU

[joshwa-s]

I am unable to run AF3x on a new blackwell A6000 GPU. I can successfully run AF3x on the old RTX A6000 (48 Gb). I can successfully run AF3 on either GPU. I tried rebuilding env/docker and upgrading CUDA, jaxlib, etc., but nothing has worked.

[gilep]

AF3x has not yet been updated on the main branch to the latest AlphaFold 3 version, which includes support for Blackwell GPUs.

For now, please install AF3x from the experimental branch: it includes the AlphaFold 3 fixes needed to support Blackwell GPUs.

https://github.com/KosinskiLab/af3x/tree/experimental