jax error when create_notebook.py is run

[BismahGhafoor]

I am trying to make the jupyter notebook but the script is failing to open the .pkl files.

This is the error I am getting:

(base) [bg171@alice-login01 FinalProject]$ python3 test5.py
I0716 10:35:35.385460 23456246728512 test5.py:126] now processing Q09472_and_P57682
Traceback (most recent call last):
File "/lustre/alice3/scratch/alice/b/bg171/FinalProject/test5.py", line 159, in
app.run(main)
File "/home/b/bg171/miniconda3/lib/python3.12/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/b/bg171/miniconda3/lib/python3.12/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/lustre/alice3/scratch/alice/b/bg171/FinalProject/test5.py", line 139, in main
pae_mtx, iptm_score = obtain_pae_and_iptm(result_subdir=result_subdir,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/lustre/alice3/scratch/alice/b/bg171/FinalProject/test5.py", line 50, in obtain_pae_and_iptm
data = pickle.load(f)
^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'jax._src.device_array'

I am using jax and jaxlib versions 0.4.21.

some guidance would be appreciated a lot.

[dingquanyu]

Hi @BismahGhafoor

It seems only jax.0.4.23 works with this import. Could you try installing jax0.4.23 ?

Yours
Dingquan

[BismahGhafoor]

Hi @dingquanyu I still got the same error even with jax 0.4.23

[dingquanyu]

Sorry I really have no idea how to solve this then. For me everything above 0.4.23 reported this error. Could you create an issue under the Jax or jaxlib repo? I'm sure there are many more discussions there.

[dingquanyu]

Meanwhile I will convert this to a discussion and hope there's someone else who solved this issue.

[jkosinski]

Could it be your pickles were generated with older versions than 0.4.21 ? I think device_array was deprecated around 0.4.1. Maybe you could install a copy of AP conda env and install:
pip install jax==0.3.25 jaxlib==0.3.25
which was the default in the older versions of AP.

@dingquanyu ideally we could remove jax from namespaces of the pkls before saving, if possible, to make pkls more compatible with future jax versions.

[BismahGhafoor]

How do i install a copy of AP conda env?

[jkosinski]

You can just install a second AlphaPulldown environment e.g. AlphaPulldownTest like this

conda create -n AlphaPulldownTest -c omnia -c bioconda -c conda-forge python==3.10 openmm==8.0 pdbfixer==1.9 kalign2 cctbx-base pytest importlib_metadata hhsuite modelcif

and then in the "Installation using pip" part, execute:

pip install jax==0.3.25 jaxlib==0.3.25

[BismahGhafoor]

The issue I'm having with that is this error:
ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.25 (from versions: 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, 0.4.30)
ERROR: No matching distribution found for jaxlib==0.3.25

[dingquanyu]

The issue I'm having with that is this error: ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.25 (from versions: 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, 0.4.30) ERROR: No matching distribution found for jaxlib==0.3.25

Hi
I think you have to add this flag after the pip install command -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlotherwise pip fails to find the versions

[BismahGhafoor]

it still didn't find the distribution for jaxlib. When I open the .html link the 0.3.25 is not there