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

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:

Start by reproducing the AttributeError after a Factory Reset in Colab, or update the local colabfold_batch installation and run the prediction block. Collect the complete ColabFold output and environment details requested in the issue; done means the Run Prediction block completes without the missing jax.xla error.