Function _reference_selected_logp

tests/test_logp.py:28–30

(logits: torch.Tensor, token_ids: torch.Tensor)

def _reference_selected_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
    log_probs = torch.log_softmax(logits.float(), dim=-1)
    return torch.gather(log_probs, dim=-1, index=token_ids.long().unsqueeze(-1)).squeeze(-1)


class TestNativeLogpOpCorrectness: