refactor: wrap multi_chain_permutation_alignment

Author: chungongyu

profold2/model/alphafold2.py

@@ -13,7 +13,7 @@
 from einops import rearrange
 
 from profold2.common import residue_constants
-from profold2.model import commons, functional
+from profold2.model import commons, folding, functional
 from profold2.model.evoformer import Evoformer
 from profold2.model.head import HeaderBuilder
 from profold2.utils import env, exists
@@ -259,9 +259,7 @@ def forward(
       if 'representations' in value:
         representations.update(value['representations'])
     if 'folding' in ret.headers:
-      batch = functional.multi_chain_permutation_alignment(
-          ret.headers['folding'], batch
-      )
+      batch = folding.multi_chain_permutation_alignment(ret.headers['folding'], batch)
     if self.training and compute_loss:
       for name, module, options in self.headers:
         if not hasattr(module, 'loss') or name not in ret.headers:

profold2/model/folding.py

@@ -210,3 +210,7 @@ def forward(self, representations, batch):
           )
 
     return outputs
+
+
+def multi_chain_permutation_alignment(value, batch):
+  return functional.multi_chain_permutation_alignment(value, batch)