AF3 AtomTransformerBlock类解读

AtomTransformerBlock 是一个用于处理原子特征和原子对关系的块结构模块，可能用于化学、分子建模或其他涉及原子级别的深度学习任务。其核心功能包括：

1. 基于注意力机制的原子间交互计算（AtomAttentionPairBias）。
2. 原子特征的非线性变换（ConditionedTransitionBlock）。

源代码：

class AtomTransformerBlock(nn.Module):
    def __init__(
            self,
            c_atom: int,
            c_atompair: int = 16,
            no_heads: int = 8,
            dropout=0.0,
            n_queries: int = 32,
            n_keys: int = 128,
    ):
        """Initialize a block within AtomTransformer module.
        Args:
            c_atom:
                Total dimension of the model.
            c_atompair:
                The number of channels for the atom pair representation. Defaults to 16.
            no_heads:
                Number of parallel attention heads. Note that c_atom will be split across no_heads
                (i.e. each head will have dimension c_atom // no_heads).
            dropout:
                Dropout probability on attn_output_weights. Default: 0.0 (no dropout).
            n_queries:
                The size of the atom window. Defaults to 32.