import_jax_weights_ failed on AlphaFold-Multimer 2.3.0

[dingquanyu]

Hi,

Thanks for the conversion script. It worked on AlphaFold Multimer 2.2.0 but when I run the same command on AlphaFold 2.3.0

I guess that's because AlphaFold changed the layout of outputs in 2.3.0? Do you have plans to update the convert script later on?

Traceback (most recent call last):
  File "convert_alphafold_to_unifold.py", line 17, in <module>
    import_jax_weights_(model, load_ckpt, version=model_name)
  File "/g/kosinski/geoffrey/unifold_review/unifold_classification/Uni-Fold/scripts/translate_jax_params.py", line 532, in import_jax_weights_
    assign(flat, data)
  File "/g/kosinski/geoffrey/unifold_review/unifold_classification/Uni-Fold/scripts/translate_jax_params.py", line 113, in assign
    weights = torch.as_tensor(orig_weights[k])
  File "/home/dingquanyu/.conda/envs/unifold_3.8/lib/python3.8/site-packages/numpy/lib/npyio.py", line 260, in __getitem__
    raise KeyError("%s is not a file in the archive" % key)
KeyError: 'alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input//scale is not a file in the archive'

I have attached the full error report as well: here and look forward to your advice. Thank you.