Embedding Bag Sum

Problem

Sum embedding rows per bag. Bag k covers indices from offsets[k] up to offsets[k + 1], or to the end for the last bag. Return one summed vector per bag.

The signature is embedding_bag_sum(indices, offsets, table).

Hint

Offsets mark bag starts; sliced ranges are contiguous in the indices list.

Test cases

case 1

in: [[0,1,2,1],[0,2],[[1,2],[3,4],[5,6]]]

expected: [[4,6],[8,10]]

case 2

in: [[0,0],[0],[[1,2]]]

expected: [[2,4]]

case 3

in: [[1,0,1],[0,1],[[1,0],[0,1]]]

expected: [[0,1],[1,1]]

case 4

in: [[2,2,2],[0,1,2],[[1],[2],[3]]]

expected: [[3],[3],[3]]

Starter codePython

def embedding_bag_sum(indices, offsets, table):
    # Your code here
    pass

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