Issue with 'jax' module when opening .pkl files in ChimeraX

[sritharan@…]

Dear ChimeraX Support,

I am currently experiencing an issue with ChimeraX that I am hoping you can help me with. When I attempt to open .pkl files using ChimeraX, I receive the following error message:

arduino

ModuleNotFoundError: No module named 'jax'
File "/usr/libexec/UCSF-ChimeraX-daily/lib/python3.9/site-packages/chimerax/alphafold/pae.py", line 887, in read_pickle_pae_matrix
p = pickle.load(f)

This error message suggests that the 'jax' module, which is a dependency for these .pkl files, is not found.

To try to resolve this issue, I have attempted to install 'jax' and 'jaxlib' using pip, specifically for Python 3.9 as indicated by the error message. Despite successfully installing 'jax' and 'jaxlib' in my Python 3.9 environment (as far as I can tell), I continue to encounter the same error message in ChimeraX.

Given that the error persists even after installing 'jax' and 'jaxlib', I suspect that ChimeraX might be using a different Python environment where these modules are not installed. However, I am not sure how to configure ChimeraX to use the correct Python environment, or if this is indeed the source of the problem.

I would greatly appreciate any guidance you could provide to help me resolve this issue. Thank you very much for your assistance.

Best regards,

[pett]

Reported by Sujith Sritharan

[Tom Goddard]

You are right ChimeraX has its own Python environment. But you don't need to add jax. This is fixed in the current ChimeraX version 1.6.1. The problem was a bug in AlphaFold itself that accidentally put a jax dependency into the .pkl file. I reported that bug about 6 months ago and it was fixed in the github AlphaFold code several months ago. So running current AlphaFold versions should not produce these problematic .pkl files. To handle the ones people already created the new ChimeraX takes care of the jax dependency (by making a dummy jax module since the jax data in the pkl file is not used by ChimeraX).