CollabFold: module 'jax' has no attribute 'tree_multimap'

[eduardo_cheuiche_rs@…]

The following bug report has been submitted:
Platform:        Windows-10-10.0.19044
ChimeraX Version: 1.5.dev202208020103 (2022-08-02 01:03:18 UTC)
Description
Colabfold has been unable to work due to a bug in the module jax, the error massage is this one:

UnfilteredStackTrace: AttributeError: module 'jax' has no attribute 'tree_multimap'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------
The above exception was the direct cause of the following exception:

AttributeError                            Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/haiku/_src/stateful.py in difference(before, after)
    310   params_before, params_after = box_and_fill_missing(before.params,
    311                                                      after.params)
--> 312   params_after = jax.tree_multimap(functools.partial(if_changed, is_new_param),
    313                                    params_before, params_after)
    314 

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

Log:
UCSF ChimeraX version: 1.5.dev202208020103 (2022-08-02)  
© 2016-2022 Regents of the University of California. All rights reserved.  
How to cite UCSF ChimeraX  

> ui tool show AlphaFold

> alphafold predict
> EVQLVESGGGLVQAGGSLRLSCAASGRTFSSYAMGWFRQAPGKEREFVVAINWSSGSTYYADSVKGRFTISRDNAKNTMYLQMNSLKPEDTAVYYCAAGYQINSGNYNFKDYEYDYWGQGTQVTVSSALEHHHHHH
> templates true

Please cite ColabFold: Making protein folding accessible to all. Nature
Methods (2022) if you use these predictions.  
Running AlphaFold prediction  




OpenGL version: 3.3.0 NVIDIA 456.71
OpenGL renderer: GeForce GTX 760/PCIe/SSE2
OpenGL vendor: NVIDIA Corporation

Python: 3.9.11
Locale: pt_BR.cp1252
Qt version: PyQt6 6.3.1, Qt 6.3.1
Qt runtime version: 6.3.1
Qt platform: windows

Manufacturer: ASUS
Model: All Series
OS: Microsoft Windows 10 Home (Build 19044)
Memory: 17,114,587,136
MaxProcessMemory: 137,438,953,344
CPU: 8 Intel(R) Core(TM) i7-4770K CPU @ 3.50GHz
OSLanguage: pt-BR

Installed Packages:
    alabaster: 0.7.12
    appdirs: 1.4.4
    asttokens: 2.0.5
    Babel: 2.10.3
    backcall: 0.2.0
    blockdiag: 3.0.0
    build: 0.7.0
    certifi: 2022.6.15
    cftime: 1.6.1
    charset-normalizer: 2.1.0
    ChimeraX-AddCharge: 1.2.3
    ChimeraX-AddH: 2.1.11
    ChimeraX-AlignmentAlgorithms: 2.0
    ChimeraX-AlignmentHdrs: 3.2.1
    ChimeraX-AlignmentMatrices: 2.0
    ChimeraX-Alignments: 2.5.2
    ChimeraX-AlphaFold: 1.0
    ChimeraX-AltlocExplorer: 1.0.3
    ChimeraX-AmberInfo: 1.0
    ChimeraX-Arrays: 1.0
    ChimeraX-Atomic: 1.39.7
    ChimeraX-AtomicLibrary: 7.0.2
    ChimeraX-AtomSearch: 2.0.1
    ChimeraX-AxesPlanes: 2.1
    ChimeraX-BasicActions: 1.1.2
    ChimeraX-BILD: 1.0
    ChimeraX-BlastProtein: 2.1.2
    ChimeraX-BondRot: 2.0.1
    ChimeraX-BugReporter: 1.0.1
    ChimeraX-BuildStructure: 2.7.1
    ChimeraX-Bumps: 1.0
    ChimeraX-BundleBuilder: 1.2
    ChimeraX-ButtonPanel: 1.0.1
    ChimeraX-CageBuilder: 1.0.1
    ChimeraX-CellPack: 1.0
    ChimeraX-Centroids: 1.2
    ChimeraX-ChangeChains: 1.0
    ChimeraX-CheckWaters: 1.1
    ChimeraX-ChemGroup: 2.0
    ChimeraX-Clashes: 2.2.4
    ChimeraX-ColorActions: 1.0.1
    ChimeraX-ColorGlobe: 1.0
    ChimeraX-ColorKey: 1.5.2
    ChimeraX-CommandLine: 1.2.4
    ChimeraX-ConnectStructure: 2.0.1
    ChimeraX-Contacts: 1.0.1
    ChimeraX-Core: 1.5.dev202208020103
    ChimeraX-CoreFormats: 1.1
    ChimeraX-coulombic: 1.3.2
    ChimeraX-Crosslinks: 1.0
    ChimeraX-Crystal: 1.0
    ChimeraX-CrystalContacts: 1.0.1
    ChimeraX-DataFormats: 1.2.2
    ChimeraX-Dicom: 1.1
    ChimeraX-DistMonitor: 1.1.6
    ChimeraX-Dssp: 2.0
    ChimeraX-EMDB-SFF: 1.0
    ChimeraX-ExperimentalCommands: 1.0
    ChimeraX-FileHistory: 1.0.1
    ChimeraX-FunctionKey: 1.0.1
    ChimeraX-Geometry: 1.2
    ChimeraX-gltf: 1.0
    ChimeraX-Graphics: 1.1
    ChimeraX-Hbonds: 2.2.1
    ChimeraX-Help: 1.2.1
    ChimeraX-HKCage: 1.3
    ChimeraX-IHM: 1.1
    ChimeraX-ImageFormats: 1.2
    ChimeraX-IMOD: 1.0
    ChimeraX-IO: 1.0.1
    ChimeraX-ItemsInspection: 1.0.1
    ChimeraX-Label: 1.1.5
    ChimeraX-ListInfo: 1.1.1
    ChimeraX-Log: 1.1.5
    ChimeraX-LookingGlass: 1.1
    ChimeraX-Maestro: 1.8.1
    ChimeraX-Map: 1.1.1
    ChimeraX-MapData: 2.0
    ChimeraX-MapEraser: 1.0.1
    ChimeraX-MapFilter: 2.0
    ChimeraX-MapFit: 2.0
    ChimeraX-MapSeries: 2.1.1
    ChimeraX-Markers: 1.0.1
    ChimeraX-Mask: 1.0.1
    ChimeraX-MatchMaker: 2.0.6
    ChimeraX-MDcrds: 2.6
    ChimeraX-MedicalToolbar: 1.0.2
    ChimeraX-Meeting: 1.0.1
    ChimeraX-MLP: 1.1
    ChimeraX-mmCIF: 2.7
    ChimeraX-MMTF: 2.1
    ChimeraX-Modeller: 1.5.6
    ChimeraX-ModelPanel: 1.3.6
    ChimeraX-ModelSeries: 1.0.1
    ChimeraX-Mol2: 2.0
    ChimeraX-Morph: 1.0
    ChimeraX-MouseModes: 1.1.1
    ChimeraX-Movie: 1.0
    ChimeraX-Neuron: 1.0
    ChimeraX-Nucleotides: 2.0.3
    ChimeraX-OpenCommand: 1.9.1
    ChimeraX-PDB: 2.6.7
    ChimeraX-PDBBio: 1.0
    ChimeraX-PDBLibrary: 1.0.2
    ChimeraX-PDBMatrices: 1.0
    ChimeraX-PickBlobs: 1.0.1
    ChimeraX-Positions: 1.0
    ChimeraX-PresetMgr: 1.1
    ChimeraX-PubChem: 2.1
    ChimeraX-ReadPbonds: 1.0.1
    ChimeraX-Registration: 1.1.1
    ChimeraX-RemoteControl: 1.0
    ChimeraX-RenumberResidues: 1.1
    ChimeraX-ResidueFit: 1.0.1
    ChimeraX-RestServer: 1.1
    ChimeraX-RNALayout: 1.0
    ChimeraX-RotamerLibMgr: 2.0.1
    ChimeraX-RotamerLibsDunbrack: 2.0
    ChimeraX-RotamerLibsDynameomics: 2.0
    ChimeraX-RotamerLibsRichardson: 2.0
    ChimeraX-SaveCommand: 1.5.1
    ChimeraX-SchemeMgr: 1.0
    ChimeraX-SDF: 2.0
    ChimeraX-Segger: 1.0
    ChimeraX-Segment: 1.0
    ChimeraX-SelInspector: 1.0
    ChimeraX-SeqView: 2.7.2
    ChimeraX-Shape: 1.0.1
    ChimeraX-Shell: 1.0.1
    ChimeraX-Shortcuts: 1.1.1
    ChimeraX-ShowSequences: 1.0.1
    ChimeraX-SideView: 1.0.1
    ChimeraX-Smiles: 2.1
    ChimeraX-SmoothLines: 1.0
    ChimeraX-SpaceNavigator: 1.0
    ChimeraX-StdCommands: 1.8
    ChimeraX-STL: 1.0
    ChimeraX-Storm: 1.0
    ChimeraX-StructMeasure: 1.0.2
    ChimeraX-Struts: 1.0.1
    ChimeraX-Surface: 1.0
    ChimeraX-SwapAA: 2.0.1
    ChimeraX-SwapRes: 2.1.2
    ChimeraX-TapeMeasure: 1.0
    ChimeraX-Test: 1.0
    ChimeraX-Toolbar: 1.1.2
    ChimeraX-ToolshedUtils: 1.2.1
    ChimeraX-Tug: 1.0.1
    ChimeraX-UI: 1.21.4
    ChimeraX-uniprot: 2.2.1
    ChimeraX-UnitCell: 1.0.1
    ChimeraX-ViewDockX: 1.1.3
    ChimeraX-VIPERdb: 1.0
    ChimeraX-Vive: 1.1
    ChimeraX-VolumeMenu: 1.0.1
    ChimeraX-VTK: 1.0
    ChimeraX-WavefrontOBJ: 1.0
    ChimeraX-WebCam: 1.0.1
    ChimeraX-WebServices: 1.1.0
    ChimeraX-Zone: 1.0.1
    colorama: 0.4.5
    comtypes: 1.1.10
    cxservices: 1.2
    cycler: 0.11.0
    Cython: 0.29.32
    debugpy: 1.6.2
    decorator: 5.1.1
    docutils: 0.19
    entrypoints: 0.4
    executing: 0.9.1
    filelock: 3.4.2
    fonttools: 4.34.4
    funcparserlib: 1.0.0
    grako: 3.16.5
    h5py: 3.7.0
    html2text: 2020.1.16
    idna: 3.3
    ihm: 0.33
    imagecodecs: 2022.7.31
    imagesize: 1.4.1
    importlib-metadata: 4.12.0
    ipykernel: 6.15.1
    ipython: 8.4.0
    ipython-genutils: 0.2.0
    jedi: 0.18.1
    Jinja2: 3.1.2
    jupyter-client: 7.3.4
    jupyter-core: 4.11.1
    kiwisolver: 1.4.4
    line-profiler: 3.4.0
    lxml: 4.9.1
    lz4: 4.0.2
    MarkupSafe: 2.1.1
    matplotlib: 3.5.2
    matplotlib-inline: 0.1.3
    msgpack: 1.0.4
    nest-asyncio: 1.5.5
    netCDF4: 1.6.0
    networkx: 2.8.5
    numexpr: 2.8.3
    numpy: 1.23.1
    openvr: 1.16.802
    packaging: 21.3
    ParmEd: 3.4.3
    parso: 0.8.3
    pep517: 0.12.0
    pickleshare: 0.7.5
    Pillow: 9.2.0
    pip: 21.3.1
    pkginfo: 1.8.2
    prompt-toolkit: 3.0.30
    psutil: 5.9.1
    pure-eval: 0.2.2
    pycollada: 0.7.2
    pydicom: 2.3.0
    Pygments: 2.12.0
    PyOpenGL: 3.1.5
    PyOpenGL-accelerate: 3.1.5
    pyparsing: 3.0.9
    PyQt6-commercial: 6.3.1
    PyQt6-Qt6: 6.3.1
    PyQt6-sip: 13.4.0
    PyQt6-WebEngine-commercial: 6.3.1
    PyQt6-WebEngine-Qt6: 6.3.1
    python-dateutil: 2.8.2
    pytz: 2022.1
    pywin32: 303
    pyzmq: 23.2.0
    qtconsole: 5.3.1
    QtPy: 2.1.0
    RandomWords: 0.3.0
    requests: 2.28.1
    scipy: 1.9.0
    setuptools: 62.6.0
    sfftk-rw: 0.7.2
    six: 1.16.0
    snowballstemmer: 2.2.0
    sortedcontainers: 2.4.0
    Sphinx: 5.1.1
    sphinx-autodoc-typehints: 1.19.1
    sphinxcontrib-applehelp: 1.0.2
    sphinxcontrib-blockdiag: 3.0.0
    sphinxcontrib-devhelp: 1.0.2
    sphinxcontrib-htmlhelp: 2.0.0
    sphinxcontrib-jsmath: 1.0.1
    sphinxcontrib-qthelp: 1.0.3
    sphinxcontrib-serializinghtml: 1.1.5
    stack-data: 0.3.0
    tables: 3.7.0
    tifffile: 2022.7.31
    tinyarray: 1.2.4
    tomli: 2.0.1
    tornado: 6.2
    traitlets: 5.3.0
    urllib3: 1.26.11
    wcwidth: 0.2.5
    webcolors: 1.12
    wheel: 0.37.1
    wheel-filename: 1.3.0
    WMI: 1.5.1
    zipp: 3.8.1

[Eric Pettersen]

Reported by Eduardo C Antonio

[Tom Goddard]

I see if you use the PDB templates in the alphafold prediction it hits this error due to a jax update that eliminated a function used by the older haiku module that ColabFold is using. It works with some warnings if you do not use PDB structure templates.

Probably the best way for me to fix this is to update the ColabFold version that ChimeraX uses. But I have to test that extensively first. Another choice is I might make ColabFold use an older version of jax that worked.

[Tom Goddard]

Fixed. No need to update ChimeraX because fix is in server script so older ChimeraX will use it.

Updated ColabFold to latest version which uses a newer dm-haiku module that replaces the deprecated jax call. It is disturbing to see that ColabFold installs most of its dependencies without using pinned versions allowing newer patch releases. And for jax it allows any newer versions of 0.3.x and the release of a new jax is what caused this bug. Ideally ChimeraX would use a ColabFold with exact versions of all dependencies specified so it won't break when updates occur. I've made a separate ticket #7762 to see if that can be done.