Function batched_sample_onehot

flowdock/utils/model_utils.py:207–216

Implements the Gumbel-Max trick to sample from a one-hot distribution.

(logits, dim=0, max_only=False)

def batched_sample_onehot(logits, dim=0, max_only=False):
    """Implements the Gumbel-Max trick to sample from a one-hot distribution."""
    if max_only:
        sampled_idx = torch.argmax(logits, dim=dim, keepdim=True)
    else:
        noise = torch.rand_like(logits)
        sampled_idx = torch.argmax(logits - torch.log(-torch.log(noise)), dim=dim, keepdim=True)
    out_onehot = torch.zeros_like(logits, dtype=torch.bool)
    out_onehot.scatter_(dim=dim, index=sampled_idx, value=1)
    return out_onehot


def topk_edge_mask_from_logits(scores, k, randomize=False):