[BUG] function get_sparse_attention_indices_with_inter_chain in RFD3

Hi all,

I was using the code with LOW_MEMOY_MODE activated when I think I have spotted a bug in how the attention indices for the inter-chain atoms are determined. As a context, I am generating a protein with Icosahedral symmetry where each Asymmetric Subunit (chain) has 40 AA. There are a total of 60 symmetrical chains in the Icosahedral symmetry

Currently, the function get_sparse_attention_indices_with_inter_chain inside models/rfd3/src/rfd3/model/layers/block_utils.py does

for c in unique_chains: 
    query_chain = chain_id[c]
    other_chain_mask = (chain_id != query_chain) & base_mask[c, :]

So the bug is:

• Loop iterates over chain IDs (0, 1, 2, ..., 59)
• Uses those as atom indices (chain_id[c], base_mask[c, :])
• Only atoms 0-59 get inter-chain neighbors computed
• Atoms 60-33599 are never processed by this loop, so they get zeros in their inter-chain slots

Therefore, I think the loop should iterate over all query atoms (0 to L-1), not over chain IDs.

My proposed fix for this whole function is the following:

@torch.no_grad()
def get_sparse_attention_indices_with_inter_chain(
    tok_idx, D_LL, n_seq_neighbours, k_intra, k_inter, chain_id, base_mask
):
    """
    Create attention indices that guarantee inter-chain interactions for clash avoidance.

    Args:
        tok_idx: atom to token mapping [L]
        D_LL: pairwise distances [B, L, L]
        n_seq_neighbours: number of sequence neighbors
        k_intra: number of intra-chain attention keys
        k_inter: number of inter-chain attention keys
        chain_id: chain IDs for each atom [L]
        base_mask: base mask for valid pairs [L, L]

    Returns:
        attn_indices: [B, L, k_total] where k_total = k_intra + k_inter
    """
    B, L, _ = D_LL.shape
    device = D_LL.device

    # Get regular intra-chain indices (limited to k_intra)
    intra_indices = get_sparse_attention_indices(
        tok_idx, D_LL, n_seq_neighbours, k_intra, chain_id, base_mask
    )  # [B, L, k_intra]

    # Get inter-chain indices for clash avoidance - VECTORIZED
    inter_indices = torch.zeros(B, L, k_inter, dtype=torch.long, device=device)

    for b in range(B):
        # Build inter-chain mask for ALL query atoms at once: [L, L]
        # other_chain_mask[q, j] = True if chain_id[j] != chain_id[q] AND base_mask[q, j]
        same_chain = chain_id[:, None] == chain_id[None, :]  # [L, L]
        other_chain_mask = ~same_chain & base_mask  # [L, L]

        # Set distances for same-chain (and disallowed) pairs to inf
        D_inter = D_LL[b].clone()  # [L, L]
        D_inter[~other_chain_mask] = float('inf')

        # Get k_inter closest atoms from other chains for ALL atoms
        # torch.topk with largest=False gets the k smallest values
        _, inter_idx = torch.topk(D_inter, k_inter, dim=-1, largest=False)  # [L, k_inter]
        
        inter_indices[b] = inter_idx

    # Combine intra and inter chain indices
    combined_indices = torch.cat(
        [intra_indices, inter_indices], dim=-1
    )  # [B, L, k_total]

    return combined_indices

It would be great if someone could have a look at this and let me know what they think!

Start in models/rfd3/src/rfd3/model/layers/block_utils.py at get_sparse_attention_indices_with_inter_chain and trace how chain_id, base_mask, and D_LL define the returned indices. Reproduce the low-memory icosahedral case with 60 chains of 40 residues each, then verify that inter-chain slots are populated for all query atoms rather than only the first chain IDs.