Arbitrary-order autoregressive graph mask

Your task

Function signature

permuted_autoregressive_mask(
    permutation: torch.Tensor,
    neighbors: torch.Tensor,
    node_mask: torch.Tensor,
    design_mask: torch.Tensor,
    edge_mask: torch.Tensor,
    fixed_first: bool = True,
) -> dict[str, torch.Tensor]

Worked example

Decode residues in order 2, 0, 3, 1

import torch

permutation = torch.tensor([[2, 0, 3, 1]], dtype=torch.int64)
neighbors = torch.tensor([[[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]]], dtype=torch.int64)
node_mask = torch.tensor([[True, True, True, True]], dtype=torch.bool)
design_mask = torch.tensor([[True, True, True, False]], dtype=torch.bool)
edge_mask = torch.full((1, 4, 4), True, dtype=torch.bool)
fixed_first = False

result = permuted_autoregressive_mask(permutation=permutation, neighbors=neighbors, node_mask=node_mask, design_mask=design_mask, edge_mask=edge_mask, fixed_first=fixed_first)

Expected return

{
    'rank': torch.tensor([[1, 3, 0, 2]], dtype=torch.int64),
    'backward': torch.tensor([[[[False], [False], [True], [False]], [[True], [False], [True], [True]], [[False], [False], [False], [False]], [[True], [False], [True], [False]]]], dtype=torch.bool),
    'forward': torch.tensor([[[[True], [True], [False], [True]], [[False], [True], [False], [False]], [[True], [True], [True], [True]], [[False], [True], [False], [True]]]], dtype=torch.bool),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Build directed graph masks for arbitrary-order sequence decoding. permutation [B,L] is integer and each row contains every residue ID 0..L-1 exactly once, EARLIEST to LATEST. neighbors [B,L,K] gives source residue j for each target i; indices are in [0,L), B,L\ge 1, K\ge 0.
• node_mask [B,L], design_mask [B,L], edge_mask [B,L,K] are bool. True means valid residue, residue to design (unknown sequence), and enabled edge slot, respectively. Masks are selection data, not differentiable weights. Config fixed_first is bool, default True, and is the only option.
• If False, use the supplied order unchanged; design_mask does not alter it. If True, stably partition the supplied order into VALID FIXED residues (node_mask and not design_mask), then VALID DESIGNED residues, then invalid residues; preserve the original order within each group.
• Return rank [B,L], the inverse permutation of this effective schedule, including ranks for invalid residues. No random ordering or argsort is required: inverse ranks can be built from permutation equality, and stable partition ranks by pairwise comparisons.
• For a target i/source j edge, \mathrm{active} = \mathrm{edge\_mask}[i,k] AND node_mask [i] AND node_mask [j]. \mathrm{backward}[i,k,0] = \mathrm{active} AND \mathrm{rank}[j]<\mathrm{rank}[i]. \mathrm{forward}[i,k,0] = \mathrm{active} AND \mathrm{rank}[j]\ge \mathrm{rank}[i].
• Return these BOOL [B,L,K,1] masks. Self edges are strictly excluded from backward and included in forward when active. The masks are disjoint and their union equals the active edges; inactive slots are false in both. K=0 returns correctly shaped empty masks.
• Backward enables earlier SEQUENCE information only; forward enables STRUCTURE-ONLY information and must never be multiplied by a future sequence embedding. This primitive constructs masks, not a decoder or feature tensor, and therefore cannot make an already-leaking hidden representation causal.
• When fixed_first=True, all valid fixed residues precede all designed residues. When False, even fixed residues obey the supplied schedule.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def permuted_autoregressive_mask(
    permutation: torch.Tensor,  # (B, L)
    neighbors: torch.Tensor,  # (B, L, K)
    node_mask: torch.Tensor,  # (B, L)
    design_mask: torch.Tensor,  # (B, L)
    edge_mask: torch.Tensor,  # (B, L, K)
    fixed_first: bool = True,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    rank: (B, L).
    backward: (B, L, K, 1).
    forward: (B, L, K, 1)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement permuted_autoregressive_mask")

Sources and further reading

• dauparas/ProteinMPNN · protein_mpnn_utils.py
• dauparas/LigandMPNN · model_utils.py