Neural Networks & LLM Components · Easy

PyTorch Final Logits Selection

Select each padded sequence's last valid vocabulary-logit row for autoregressive decoding.

Task

Implement final_logits_selection(logits: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor. Given batched token logits with shape [B, T, V] and valid sequence lengths with shape [B], return the logits at each row's last valid token.

Requirements

logits is a rank-3 tensor with shape batch by time by vocabulary.
lengths is a rank-1 tensor with one positive valid-token count per batch row.
For row b, select time index lengths[b] - 1.
Return a tensor with shape batch by vocabulary.
Do not use logits[:, -1, :] unless every sequence has full length.
Keep the operation batched and preserve logits dtype and device.

Example

logits = torch.tensor([
    [[1.0, 2.0], [3.0, 4.0], [9.0, 9.0]],
    [[5.0, 6.0], [7.0, 8.0], [0.0, 0.0]],
])
lengths = torch.tensor([2, 1])

final_logits_selection(logits, lengths)
# tensor([[3.0, 4.0], [5.0, 6.0]])