Jax error AF2 batch notebook.

Expected Behavior

run batch predictions from google drive - used to work just fine.

Current Behavior

When running now getting this error:

AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'
Seems like something is changed in the JAX version or so ?
Could reproduce error on 3 different google accounts.

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 with the batch notebook and the colabfold_batch entry point mentioned in the report. Reproduce the error after a Colab Factory Reset, recording the complete output and environment details, then investigate the reported JAX/JAXlib mismatch. Done means batch predictions from Google Drive run successfully and the reproduction details are documented.