Fix jax compilation error in unpaired mode.

Author: dohyun-s

colabfold/alphafold/msa.py

@@ -4,6 +4,7 @@
 
 from alphafold.model.features import FeatureDict
 from alphafold.model.tf import shape_placeholders
+import jax.numpy as jnp
 
 NUM_RES = shape_placeholders.NUM_RES
 NUM_MSA_SEQ = shape_placeholders.NUM_MSA_SEQ
@@ -51,7 +52,7 @@ def make_fixed_size_multimer(
   NUM_RES = "num residues placeholder"
   NUM_MSA_SEQ = "msa placeholder"
   NUM_TEMPLATES = "num templates placeholder"
-  msa_cluster_size = feat["bert_mask"].shape[0]
+  msa_cluster_size = jnp.array(feat["bert_mask"]).shape[0]
   pad_size_map = {
     NUM_RES: num_res,
     NUM_MSA_SEQ: msa_cluster_size,

colabfold/batch.py

@@ -319,7 +319,7 @@ def main():
         host_url=args.host_url,
       )
 
-      if args.pair_mode == "none" or args.pair_mode == "unpaired" or args.pair_mode == "unpaired_paired":
+      if args.pair_mode == "none" or "unpaired" or "unpaired_paired":
         unpaired_path = Path(args.results).joinpath(str(jobname)+"_unpaired_env")
         unpaired_a3m_lines = run_mmseqs2(
           query_seqs_unique,

colabfold/predict.py

@@ -69,8 +69,8 @@ def predict_structure(
           input_features = feature_dict
           input_features["asym_id"] = input_features["asym_id"] - input_features["asym_id"][...,0]
           # TODO
+          input_features = pad_input_multimer(input_features, model_runner, model_name, pad_len,  use_templates)
           if seq_len < pad_len:
-            input_features = pad_input_multimer(input_features, model_runner, model_name, pad_len, use_templates)
             logger.info(f"Padding length to {pad_len}")
       else:
         if model_num == 0: