AF3 PairStack类源码解读

PairStack 是 AlphaFold 的核心模块之一，用于对残基对（residue-residue pair）的特征张量 z 进行迭代更新。这个模块结合几何操作（如三角形乘法）和注意力机制，逐步建模蛋白质序列中残基之间的复杂关系。

源代码：

class PairStack(nn.Module):
    def __init__(
            self,
            c_z: int,
            c_hidden_tri_mul: int = 128,
            c_hidden_pair_attn: int = 32,
            no_heads_tri_attn: int = 4,
            transition_n: int = 4,
            pair_dropout: float = 0.25,
            fuse_projection_weights: bool = False,
            inf: float = 1e8,
    ):
        super(PairStack, self).__init__()

        if fuse_projection_weights:
            self.tri_mul_out = FusedTriangleMultiplicationOutgoing(
                c_z,
                c_hidden_tri_mul,
            )
            self.tri_mul_in = FusedTriangleMultiplicationIncoming(
                c_z,
                c_hidden_tri_mul,
            )
        else:
            self.tri_mul_out = TriangleMultiplicationOutgoing(
                c_z,
                c_hidden_tri_mul,
            )
            self.tri_mul_in = TriangleMultiplicationIncoming(
                c_z,
                c_hidden_tri_mul,
            )

        self.tri_att_start = TriangleAttentionStartingNode(
            c_z,
            c_hidden_pair_attn,
            no_heads_tri_attn,
            inf=inf,
        )
        self.tri_att_end = TriangleAttentionEndingNode(
            c_z,
            c_hidden_pair_attn,
            no_heads_tri_attn,
            inf=inf,
        )

        self.transition = Transition(
            c_z,
            transition_n,
        )

        self.dropout_row_layer = DropoutRowwise(pair_dropout)
        self.dropout_col_layer = DropoutColumnwise(pair_dropout)

    def forward(