Nvidia RTX 6000 Blackwell / SM120 incompatibility

[kpetrov]

Hello,

When trying to run Colabfold on RTX6000, it crashes as the version of jax you are using is incompatible with shaders version 12.

Alphafold3 resolved this issue in the head branch by updating jax/jaxlib to 0.8.x

Do you think it would be possible to fix this in Colabfold as well? We can gladly help with testing, if needed.

thanks!
k.

[milot-mirdita]

I had briefly tried using 0.8 but ran into an issue and didn't investigate further.
Did you try to force an upgrade?

[kpetrov]

Yes, and got a whole lot of incompatibility errors. Jax seems to be evolving fast; it is at 0.9 now.

[milot-mirdita]

I just tried to run it on a dgx spark with jax 0.9.1 and it seems to run correctly. It was probably an dm-haiku issue before.