Thanks to tom, I found out that scatter_add exists, and from there I found torch_scatter’s segment_coo
Here’s my solution now:
    initial_reps_list = []
    for i, sample_output in enumerate(sequence_output):
        token_mapping = all_token_mapping[i]
        token_mapping = token_mapping[token_mapping != -1]
        non_padded_outputs = sample_output[:num_bert_tokens[i]]
        initial_reps_list.append(torch_scatter.segment_coo(non_padded_outputs, token_mapping, reduce="mean"))
    initial_reps = torch.cat(initial_reps_list)

token_mapping is a list of indices in ascending order up to the max sequence length, padded with -1. I loop through the batch, for each sample, I get the token mapping, and only keep the non-negative indices.
num_bert_tokens is a list that holds, for each sample, the number of tokens (no padding). I get the non-padded outputs, use segment_coo to reduce them according to the token_mapping, and append them all to a list.
After the loop, I concatenate all the tensors in the list together.
The method segment_coo reduces all values from the src tensor into out at the indices specified in the index tensor along the last dimension of index. More details can be found at: Segment COO — pytorch_scatter 2.0.6 documentation
Code runs much faster now