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

[hanabors]

Expected Behavior

Run Prediction block will complete

Current Behavior

Run prediction block is currently erroring with AttributeError: module 'jax' has no attribute 'xla'

Steps to Reproduce (for bugs)

Please make sure to reproduce the issue after a "Factory Reset" in Colab. If running locally ypdate you local installation colabfold_batch to the newest version. Please provide your input if you can share it.

ColabFold Output (for bugs)

Please make sure to also post the complete ColabFold output. You can use gist.github.com for large output.

Context

Providing context helps us come up with a solution and improve our documentation for the future.

Your Environment

Include as many relevant details about the environment you experienced the bug in.

• Git commit used
• If you run it on a local system. Please add the server specifications
• Operating system and version:

[sokrypton]

This issue should now be resolved. Can you try again?

Make sure you refresh the notebook (or make a copy of our official notebook, if you are using your own).

[HakonOnco]

I am running in Google Colab, and am still seeing the issue

module 'jax' has no attribute 'xla'

I have been running the exact same script for a couple of weeks, but only got the error today. I have tried to completely restart and detach and re-attach from colab, but unless there is some method of restarting I dont know about, the error is still there.

[martin-steinegger]

Could you please post the link to the notebook you are using?

[HakonOnco]

Sure. It is not exactly optimized, with some commented out sections, but it is essentially a version of one of your templates, made to connect with google drive and get fastas from there. It also stores them there afterwards

https://colab.research.google.com/drive/1ahi0ADXiXRn0dUuNU0AOKFTsZZ5KWz9D#scrollTo=AzIKiDiCaHAn

I would appreciate any help you can give me.

[sokrypton]

The issue is coming from dm-haiku, the version we have pinned is no longer compatible with the default jax version installed in colab. After this line:

os.system("pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'")

add:

os.system("pip install --upgrade dm-haiku")

[HakonOnco]

Thank you very much, it seems to be working now.