AttributeError: module 'jax' has no attribute 'Array'

[seanwangsalad]

Dear developers,

I am getting this issue when running T-fold:

AttributeError: module 'jax' has no attribute 'Array'

I am using Alphafold 2.1.0. My jax versions are: jax 0.2.24 jaxlib 0.1.69+cuda111

Full error: making Kd predictions for 1 pmhcs... /home/mhc/miniforge3/envs/tfold-env/lib/python3.8/site-packages/tensorflow/python/keras/engine/functional.py:588: UserWarning: Input dict contained keys ['regmask'] which did not match any model input. They will be ignored by the model. warnings.warn( total AF models to be produced: 15 2024-02-08 12:59:22.829698: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0 Traceback (most recent call last): File "tfold_run_alphafold.py", line 39, in from alphafold.model import model File "/choncc/alphafold/alphafold/model/model.py", line 21, in from alphafold.model import modules File "/choncc/alphafold/alphafold/model/modules.py", line 21, in from alphafold.model import all_atom File "/choncc/alphafold/alphafold/model/all_atom.py", line 40, in from alphafold.model import utils File "/choncc/alphafold/alphafold/model/utils.py", line 29, in def stable_softmax(logits: jax.Array) -> jax.Array: AttributeError: module 'jax' has no attribute 'Array' Traceback (most recent call last): File "collect_results.py", line 12, in result_parse_tools.parse_results(working_dir) File "/choncc/tfold-release/tfold/modeling/result_parse_tools.py", line 65, in parse_results pmhc_id=os.listdir(ctarget_dir+'/outputs')[0] IndexError: list index out of range

Thanks!

[seanwangsalad]

The jax version downloaded from alphafold is outdated and does not have the .Array module. After upgrading Jax and Jaxlib to 0.4.13, this error was resolved.

New error has popped up: I0208 14:00:36.032410 140157135853376 tfold_run_alphafold.py:176] Using random seed 318656999041770268 for the data pipeline I0208 14:00:36.032602 140157135853376 tfold_run_alphafold.py:87] Predicting for id 0 Traceback (most recent call last): File "tfold_run_alphafold.py", line 196, in app.run(main) File "/home/mhc/miniforge3/envs/tfold-env/lib/python3.8/site-packages/absl/app.py", line 312, in run _run_main(main, args) File "/home/mhc/miniforge3/envs/tfold-env/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main sys.exit(main(argv)) File "tfold_run_alphafold.py", line 186, in main predict_structure(sequences=sequences,msas=msas,template_hits=template_hits,renumber_list=renumber_list, File "tfold_run_alphafold.py", line 92, in predict_structure feature_dict=data_pipeline.process(sequences,msas,template_hits) File "/choncc/tfold-release/tfold_patch/tfold_pipeline.py", line 135, in process msa_features=make_msa_features(msas) File "/choncc/tfold-release/tfold_patch/tfold_pipeline.py", line 113, in make_msa_features uniprot_accession_ids.append(identifiers.uniprot_accession_id.encode('utf-8')) AttributeError: 'Identifiers' object has no attribute 'uniprot_accession_id' Traceback (most recent call last): File "collect_results.py", line 12, in result_parse_tools.parse_results(working_dir) File "/choncc/tfold-release/tfold/modeling/result_parse_tools.py", line 74, in parse_results if get_rmsd: UnboundLocalError: local variable 'get_rmsd' referenced before assignment

[v-mikhaylov]

Hi Sean, thank you for your interest in our algorithm! For me it works with AF 2.1.0 and jax 0.2, certainly not jax 0.4. I looked into the code of AF2.1.0 (downloaded as a zip file from https://github.com/google-deepmind/alphafold/releases/tag/v2.1.0) and alphafold.model.utils does not have a function stable_softmax. On the other hand, AF2.3 does have this function in alphafold.model.utils and the function indeed uses jax.Array which I believe was added to jax in version 0.4. Can it be that you are actually using AF version >2.1.0?

[seanwangsalad]

Yes, it seems I was on AF 2.3. My bad!