TemplatePairStack是实现蛋白质结构模版pair_act特征表示的类:
通过layer_stack.layer_stack(c.num_block)(block) 堆叠c.num_block(配置文件中为2)block 函数,每个block对输入pair_act 和 pair_mask执行计算流程:TriangleAttention —> dropout ->TriangleAttention —> dropout -> TriangleMultiplication —> dropout -> TriangleMultiplication —> dropout -> Transition
import haiku as hkclass TemplatePairStack(hk.Module):"""Pair stack for the templates.Jumper et al. (2021) Suppl. Alg. 16 "TemplatePairStack""""def __init__(self, config, global_config, name='template_pair_stack'):super().__init__(name=name)self.config = configself.global_config = global_configdef __call__(self, pair_act, pair_mask, is_training, safe_key=None):"""Builds TemplatePairStack module.Arguments:pair_act: Pair activations for single template,