Catch the issue of not being able to find ptxas

[xiki-tempula]

If one just do a fresh conda install of pymbar 4 on a machine without cuda, one could get

│ ╭───────────────────────────────────────── locals ──────────────────────────────────────────╮    │
│ │ continuation = None                                                                       │    │
│ │  f_k_nonzero = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., │    │
│ │                │      0., 0., 0., 0., 0., 0., 0., 0., 0.])                                │    │
│ │       method = 'adaptive'                                                                 │    │
│ │   N_k_nonzeo = array([4584., 4982., 4251., 4232., 2857., 3734., 4141., 4565., 4730.,      │    │
│ │                │      4533., 4373., 4278., 5001., 4982., 4870., 4704., 4044., 3583.,      │    │
│ │                │      4179., 4490., 4437., 4616., 2181., 4681., 4576., 3712.])            │    │
│ │  N_k_nonzero = array([4584., 4982., 4251., 4232., 2857., 3734., 4141., 4565., 4730.,      │    │
│ │                │      4533., 4373., 4278., 5001., 4982., 4870., 4704., 4044., 3583.,      │    │
│ │                │      4179., 4490., 4437., 4616., 2181., 4681., 4576., 3712.])            │    │
│ │      options = {'maxiter': 10000, 'verbose': False}                                       │    │
│ │          tol = 1e-12                                                                      │    │
│ │ u_kn_nonzero = array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00, ...,            │    │
│ │                │   │    1.39702014e+01,  3.21518809e+01,  3.60152730e+01],                │    │
│ │                │      [ 5.13781006e-01,  9.31665221e-02,  2.72818891e-01, ...,            │    │
│ │                │   │    1.23757798e+01,  3.05602798e+01,  3.44667669e+01],                │    │
│ │                │      [ 1.75389511e+00,  8.02655973e-01,  1.22281803e+00, ...,            │    │
│ │                │   │    1.07887631e+01,  2.89787227e+01,  3.29260834e+01],                │    │
│ │                │      ...,                                                                │    │
│ │                │      [ 1.02462922e+12,  2.36204171e+06,  4.82963911e+07, ...,            │    │
│ │                │   │   -5.68161594e+00, -1.66627325e+00, -1.31547888e+00],                │    │
│ │                │      [ 1.02462922e+12,  2.36203915e+06,  4.82963930e+07, ...,            │    │
│ │                │   │   -3.06945073e+00, -1.06177938e+00, -8.89514293e-01],                │    │
│ │                │      [ 1.02462922e+12,  2.36203706e+06,  4.82963953e+07, ...,            │    │
│ │                │   │    0.00000000e+00,  0.00000000e+00,  0.00000000e+00]])               │    │
│ ╰───────────────────────────────────────────────────────────────────────────────────────────╯    │
│                                                                                                  │
│ /home/ubuntu/mambaforge/envs/asfe/lib/python3.8/site-packages/pymbar/mbar_solvers.py:660 in      │
│ precondition_u_kn                                                                                │
│                                                                                                  │
│   657 │   x_n such that the current objective function value is zero, which                      │
│   658 │   should give maximum precision in the objective function.                               │
│   659 │   """                                                                                    │
│ ❱ 660 │   return jax_precondition_u_kn(u_kn, N_k, f_k)                                           │
│   661                                                                                            │
│   662                                                                                            │
│   663 def solve_mbar_once(                                                                       │
│                                                                                                  │
│ ╭───────────────────────────────────── locals ──────────────────────────────────────╮            │
│ │  f_k = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., │            │
│ │        │      0., 0., 0., 0., 0., 0., 0., 0., 0.])                                │            │
│ │  N_k = array([4584., 4982., 4251., 4232., 2857., 3734., 4141., 4565., 4730.,      │            │
│ │        │      4533., 4373., 4278., 5001., 4982., 4870., 4704., 4044., 3583.,      │            │
│ │        │      4179., 4490., 4437., 4616., 2181., 4681., 4576., 3712.])            │            │
│ │ u_kn = array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00, ...,            │            │
│ │        │   │    1.39702014e+01,  3.21518809e+01,  3.60152730e+01],                │            │
│ │        │      [ 5.13781006e-01,  9.31665221e-02,  2.72818891e-01, ...,            │            │
│ │        │   │    1.23757798e+01,  3.05602798e+01,  3.44667669e+01],                │            │
│ │        │      [ 1.75389511e+00,  8.02655973e-01,  1.22281803e+00, ...,            │            │
│ │        │   │    1.07887631e+01,  2.89787227e+01,  3.29260834e+01],                │            │
│ │        │      ...,                                                                │            │
│ │        │      [ 1.02462922e+12,  2.36204171e+06,  4.82963911e+07, ...,            │            │
│ │        │   │   -5.68161594e+00, -1.66627325e+00, -1.31547888e+00],                │            │
│ │        │      [ 1.02462922e+12,  2.36203915e+06,  4.82963930e+07, ...,            │            │
│ │        │   │   -3.06945073e+00, -1.06177938e+00, -8.89514293e-01],                │            │
│ │        │      [ 1.02462922e+12,  2.36203706e+06,  4.82963953e+07, ...,            │            │
│ │        │   │    0.00000000e+00,  0.00000000e+00,  0.00000000e+00]])               │            │
│ ╰───────────────────────────────────────────────────────────────────────────────────╯            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
XlaRuntimeError: FAILED_PRECONDITION: Couldn't get ptxas/nvlink version string: INTERNAL: Couldn't invoke ptxas --version

It might be good to test if ptxas is there when importing pymbar 4 and advise the use to do conda install -c nvidia cuda-nvcc if it is not there.

[mikemhenry]

Thanks for the bug report! We will need to investigate why the conda-forge package doesn't pull in the correct dependencies.