Padding Side Logits Pick

Problem

Select the logits of the last real token for each sequence in a padded batch. With right padding the last real token is at index lengths[i] - 1; with left padding it is always at index seq_len - 1.

The signature is last_token_logits(logits, lengths, pad_side). logits has shape (batch, seq_len, vocab).

Hint

Left padding lets every sequence's final logits sit at the same index for one clean gather.

Test cases

case 1

in: [[[[1,2],[3,4]],[[5,6],[7,8]]],[1,2],"right"]

expected: [[1,2],[7,8]]

case 2

in: [[[[1,2],[3,4]],[[5,6],[7,8]]],[1,2],"left"]

expected: [[3,4],[7,8]]

case 3

in: [[[[1,2],[3,4],[5,6]]],[2],"right"]

expected: [[3,4]]

case 4

in: [[[[1,2],[3,4],[5,6]]],[2],"left"]

expected: [[5,6]]

case 5

in: [[[[1,2],[3,4]],[[5,6],[7,8]]],[2,1],"right"]

expected: [[3,4],[5,6]]

Starter codePython

def last_token_logits(logits, lengths, pad_side):
    # Your code here
    pass

5 test cases · run your code against them in the editor