Constrained categorical sampling with tied positions

Your task

Function signature

constrained_categorical_sampling(
    logits: torch.Tensor,
    temperature: torch.Tensor,
    uniform: torch.Tensor,
    tokens: torch.Tensor,
    group_ids: torch.Tensor,
    allowed: torch.Tensor,
    parameters: dict,
    masks: dict,
    top_k: int = 0,
    top_p: float = 1.0,
    filter_order: str = 'before_temperature',
    nucleus_rule: str = 'standard',
) -> dict[str, torch.Tensor]

Worked example

Unequal member weights and intersected legal classes

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• logits and bias are finite floating [B,L,V], \mathrm{weights}[B,L]\ge 0, \mathrm{temperature}[B,G]\ge 0; B,G\ge 1,L\ge 0,V\ge 2. uniform [B,G] is supplied in [0,1), never drawn internally. All floats share dtype/device; all scaled scores/differences must be representable.
• group_ids [B,L] is in [0,G) on valid positions (arbitrary integers on padding). valid_mask,update_mask are bool [B,L]; allowed [B,L,V] uses True means legal class. tokens [B,L] are current integers: valid non-updated positions must have IDs in [0,V), while update/padding slots may carry arbitrary placeholders.
• For group g, members are valid positions with that group ID and active members additionally have update_mask=True. Group score z_v is the SUM over active members of \mathrm{weights\_i}\cdot (\mathrm{logits\_iv}+\mathrm{bias\_iv}), divided by the NUMBER of active members, not \sum(\mathrm{weights}).
• Thus tied all-active positions with unit weights average their logits/biases. Legal classes are the INTERSECTION of allowed sets over ALL members.
• Every fixed (valid non-updated) member additionally forces its current identity, so consistent fixed ties force a one-hot outcome; contradictory fixed identities or forbidden fixed identities leave empty support. Groups with no active members also have empty support. Config top_k is 0 (default, no top-k) or an integer1..V.
• Rank legal classes by descending score with smaller class ID breaking ties and retain k if enabled. Config top_p is in (0,1], default 1.
• If \mathrm{top\_p}<1, softmax the FILTER scores on the top-k-retained support and apply nucleus_rule: 'standard' (default) keeps each class whose cumulative mass before it is < p, including the first crossing class; 'esm3' keeps classes whose cumulative mass THROUGH them is \le p and always keeps the highest-ranked legal class, thus normally excluding the first crossing class.
• For p=1 neither rule filters. Empty support stays empty under both rules. Config filter_order is 'before_temperature' (default, filter z) or 'after_temperature' (filter z/T). top-k always precedes nucleus filtering.
• For T=0, use z for filtering. After filtering, T>0 yields \operatorname{softmax}(z/T); T=0 yields a one-hot at the highest retained z, ties to smaller ID. These are all keys.
• For T>0, sample the first class in ORIGINAL ascending vocabulary order whose cumulative final probability is STRICTLY greater than uniform. If roundoff leaves no crossing, choose the last positive-probability class.
• Return sampled [B,G] (-1 for empty support), probabilities [B,G,V], group_valid [B,G], \mathrm{log\_probability}[B,G]=\operatorname{log}(probability of sampled class) or 0 if empty, and updated \mathrm{tokens}[B,L]/\mathrm{updated}[B,L].
• Only active positions in valid groups change; invalid groups leave tokens unchanged, and all fixed/padding positions are preserved. Invalid-group probabilities are exactly zero. No in-place mutation. This is one categorical update from supplied logits, not a neural decoder or a source-specific random-number stream.

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 constrained_categorical_sampling(
    logits: torch.Tensor,  # (B, L, V)
    temperature: torch.Tensor,  # (B, G)
    uniform: torch.Tensor,  # (B, G)
    tokens: torch.Tensor,  # (B, L)
    group_ids: torch.Tensor,  # (B, L)
    allowed: torch.Tensor,  # (B, L, V)
    parameters: dict,
    masks: dict,
    top_k: int = 0,
    top_p: float = 1.0,
    filter_order: str = 'before_temperature',
    nucleus_rule: str = 'standard',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    tokens: (B, L).
    updated: (B, L).
    sampled: (B, G).
    group_valid: (B, G).
    probabilities: (B, G, V).
    log_probability: (B, G)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement constrained_categorical_sampling")

Sources and further reading

• www.science.org
• dauparas/ProteinMPNN · protein_mpnn_utils.py
• www.science.org
• Biohub/esm · sampling.py
• doi.org
• togethercomputer/stripedhyena · sample.py
• arxiv.org
• salesforce/progen · sample.py