src/transformers/models/splinter/modeling_splinter.py/QuestionAwareSpanSelectionHead/forward
class QuestionAwareSpanSelectionHead:    def forward(self, inputs, positions):
        _, _, dim = inputs.size()
        index = positions.unsqueeze(-1).repeat(1, 1, dim)  # [batch_size, num_positions, dim]
        gathered_reps = torch.gather(inputs, dim=1, index=index)  # [batch_size, num_positions, dim]

        query_start_reps = self.query_start_transform(gathered_reps)  # [batch_size, num_positions, dim]
        query_end_reps = self.query_end_transform(gathered_reps)  # [batch_size, num_positions, dim]
        start_reps = self.start_transform(inputs)  # [batch_size, seq_length, dim]
        end_reps = self.end_transform(inputs)  # [batch_size, seq_length, dim]

        hidden_states = self.start_classifier(query_start_reps)  # [batch_size, num_positions, dim]
        start_reps = start_reps.permute(0, 2, 1)  # [batch_size, dim, seq_length]
        start_logits = torch.matmul(hidden_states, start_reps)

        hidden_states = self.end_classifier(query_end_reps)
        end_reps = end_reps.permute(0, 2, 1)
        end_logits = torch.matmul(hidden_states, end_reps)

        return start_logits, end_logits