Test hangs on GPU with jaxlib 0.4.2

[bwohlberg]

In scico/test/test_solver.py, the test test_minimize (to be renamed test_minimize_vector in an upcoming PR) hangs at the line

out = solver.minimize(f, x0=snp.zeros_like(x), method=method)

when run on a GPU device with jaxlib 0.4.2. It does not hang on the same device with jaxlib 0.3.15.

[Sibgatulin]

While I recognise that scico supports jax{,lib} <= 0.4.2, I gave it a try with jax==0.4.8 and jaxlib==0.4.7+cuda11.cudnn86 and the test in question runs just fine (with a few unrelated warnings) while hanging on jaxlib 0.4.2. with the same local cuda==11.8.0 and cudnn=8.6.0). Just thought I chime in in case it may be helpful.

[bwohlberg]

Thanks for pointing this out. That certainly is useful to know - if some version of jaxlib later than 0.4.2 resolves the problem then we can resolve this issue by simply bumping the minimum supported version. (The jaxlib / jax versions you mentioned will soon be included in the supported version range -- see PR #390.)

[bwohlberg]

In fact, this seems to have been fixed as far back as 0.4.3.