Last Real Token Logits

Problem

For each sequence in a padded batch, return the vocabulary-logit row at its greatest token position whose mask is true. A sequence with no true mask entry returns a zero vector of vocabulary width.

Masks may have left padding, right padding, or gaps. Counting true tokens does not identify the last valid index in the general case. Return shape B × V.

def last_token_logits(logits, mask):

Constraints

• logits is B × T × V and mask is B × T booleans; dimensions are 1..16.
• All logits are finite.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      [
        1,
        2
      ],
      [
        3,
        4
      ],
      [
        9,
        9
      ]
    ]
  ],
  [
    [
      true,
      true,
      false
    ]
  ]
]
Expected result: [
  [
    3,
    4
  ]
]

Arguments: [
  [
    [
      [
        9,
        9
      ],
      [
        1,
        2
      ],
      [
        3,
        4
      ]
    ]
  ],
  [
    [
      false,
      true,
      true
    ]
  ]
]
Expected result: [
  [
    3,
    4
  ]
]