import torch


def per_symbol_to_per_seq_probs(per_symbol_probs, tgt_out_idx):
    """ Gather per-symbol probabilities into per-seq probabilities """
    # per_symbol_probs shape: batch_size, seq_len, candidate_size
    # tgt_out_idx shape: batch_size, seq_len
    # output shape: batch_size, 1
    return torch.clamp(
        torch.prod(
            torch.gather(per_symbol_probs, 2, tgt_out_idx.unsqueeze(2)).squeeze(2),
            dim=1,
            keepdim=True,
        ),
        # prevent zero probabilities, which cause torch.log return -inf
        min=1e-40,
    )