Module pair_init 

Pair representation initialization for the ESMFold2 FoldingTrunk.

Produces the initial pair tensor [B, N, N, d_pair] from three sources:

1. Relative position encoding — for each (i,j), the clipped relative residue-index difference clip(j-i, -n_bins, n_bins) is one-hot encoded into 2*n_bins+1 = 65 bins.
2. Chain indicator — two extra bins: same-chain / cross-chain.
3. Outer product — the single representation [B,N,d_single] is projected to d_outer, outer-producted with itself, then projected to d_pair.

The three contributions are projected to d_pair and summed.

This module contains standalone math functions (relpos_encoding, chain_pair_features, outer_product) that are testable without weights, plus the PairInit struct that holds the learned projections.

Structs

PairInit

Learnable pair-representation initializer.

Functions

chain_pair_features

Per-pair same-chain / different-chain indicator.

outer_product

Flat outer product of two [B, N, d] tensors.

relpos_encoding

Relative-residue-index one-hot encoding for every (i,j) pair.