Jax dependency

[Aroids]

It appears analyzing the output of Alphafold v2.3.1 requires Jax.

Processing ranked model 0 (result_model_5_multimer_v3_pred_2).
Traceback (most recent call last):
  File "/Users/Mine/AlphaPickle_v153/alphapickle/run_AlphaPickle.py", line 77, in <module>
    results = ap.AlphaFoldPickle(args.output_directory + "/result_" + model[1] + ".pkl", args.fasta_file, ranking=str(model[0]))
  File "/Users/Mine/AlphaPickle_v153/alphapickle/src/AlphaPickle.py", line 100, in __init__
    self.data.append(pkl.load(openfile))
ModuleNotFoundError: No module named 'jax'

I did not have this problem with Alphafold v2.2.2. It sounds like this dependency was introduced in ~Alphafold v2.2.4 and may be removed in future updates (see here).

I can resolve the issue by installing jax[cpu] within my Alphapickle conda environment, but I nevertheless thought you should know since the dependency it isn't met by Alphapickle's requirements.txt.

[TienMPhan]

Hi @Aroids,
I used AF-2.3.0 and got a similar error above. l installed jax using pip install --upgrade "jax[cpu]" (jax=0.4.13, jaxlib = 0.4.13 were installed), but I still got the error:

 line 48, in __init__
    self.data.append(pkl.load(openfile))
ModuleNotFoundError: No module named 'jax._src.device_array'

How did you resolve the issue? Any specific version of jax?

[TienMPhan]

Update:
I installed older versions of jax and jaxlib as follows
pip install jax==0.3.25 jaxlib==0.3.25
and it worked on the CPU, as shown in the warning.

WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)