jaxlib not installed reverting to CPU error

[CurtisColwell]

When I updated to v1.4 I started getting these errors while running rbfe and afe:

******* JAX 64-bit mode is now on! *******
*     JAX is now set to 64-bit mode!     *
*   This MAY cause problems with other   *
*      uses of JAX in the same code.     *
******************************************

WARNING:2025-06-10 09:01:40,857:jax._src.xla_bridge:966: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

However, it does appear to still use the GPU just fine. I installed jaxlib on my developer environment and it quite ironically throws an error and does not work:

******* JAX 64-bit mode is now on! *******
*     JAX is now set to 64-bit mode!     *
*   This MAY cause problems with other   *
*      uses of JAX in the same code.     *
******************************************

E0611 09:43:08.283950  236213 pjrt_stream_executor_client.cc:2917] Execution of replica 0 failed: INTERNAL: jaxlib/gpu/solver_handle_pool.cc:37: operation gpusolverDnCreate(&handle) failed: cuSolver internal error

[CurtisColwell]

I then uninstalled it and see that it is only used for the MBAR analysis

********* JAX NOT FOUND *********
 PyMBAR can run faster with JAX  
 But will work fine without it   
Either install with pip or conda:
      pip install pybar[jax]     
               OR                
      conda install pymbar       
*********************************

[mikemhenry]

Yes sorry this is a confusing message, it doesn't have anything to do with MD happening on the GPU but just if JAX can be accelerated by the GPU. We are working to fix this upstream choderalab/pymbar#554

[IAlibay]

@mikemhenry the failure when GPU JAX is enabled might be worth looking into too. I know I've encountered that issue locally in the past. I know @jthorton has gotten it to work, so I suspect it's some kind of package installation issue.

[jthorton]

To get this working when calculating data for the industry benchmarking I just used the environment here. The way this is run is dfferent though as I only run mbar on the extracted data and not the MD, could it be an issue with running mbar on the gpu where MD is as well? Has anyone manged to do this if not I could try and test it locally.

[IAlibay]

@jthorton if you happen to have time testing it locally would be grand - I've not had any success in the past

[atravitz]

@CurtisColwell we believe this issue is resolved by #1694, but please reopen if you have any follow up concerns!