Unordered slicing and summation?

Greetings,

I am writing in order to enquire about another feature that I cannot seem to find.

I am wondering if there is a function that could do something alike “unsorted_segment_sum” in tensorflow. I’ll attempt to illustrate the purpose of this function for clarity:

Let’s say you have an array of shape [3,10,22] so a batch of size 3 with 10 elements of size 22. And you want to sum together a number of those elements or their components.

option 1:
Let’s say we want 4 classes from the 10 elements, we would need a list containing the indices for each element that we want to group togher i.e.:[[0,2,4],[1,6,8],[3,5,9],[7]] , and for each of the 3 element of the batch slice the array and sum sequentially using those group of indices.

Option 2:
The second way would be to have something with lenght 10 and to use numbers that would represent each group i.e.: [0,1,0,2,0,2,1,3,1,2]

With either case we would use those to group together the elements with the same number and perform the summation(or whatever function we would need like max, min, mean on them). The output would be either of size [3,4,22] or [4,22] depending if there is broadcasting of the query or not (though in TF the first shape must match so the latter shape is created and permutation/transposition is a necessary workaround to preserve needed dimensions.).

This could be done sequentially but this would be time intensive and not that efficient.

Let me know what you think and if I can clear this up for you.

Thanks.

I’m not aware of an operator that does that but it should be straightforward to write one. For example, here is a quick and dirty implementation of segment_sum:

class SegmentSum(mx.gluon.Block):
    
    def __init__(self, **kwargs):
        super(SegmentSum, self).__init__(**kwargs)
    
    def forward(self, x, sids):
        
        ln = max(sids)
        
        sumargs = {}
        for i in range(ln+1):
            sumargs[i] = []

        for i in range(len(sids)):
            sumargs[sids[i]].append(x[i])

        sums = []

        for i in range(ln+1):
            sums.append(nd.ElementWiseSum(*sumargs[i]))

        ans = nd.stack(*sums)
        
        return ans

You can use like like:

ss = SegmentSum()

x = nd.array([[1,2,3],[-1,-2,-3],[5,6,7]])
sids = [0,0,1]

y = ss(x, sids)
print(y)

It should return

[[ 0.  0.  0.]
 [ 5.  6.  7.]]
<NDArray 2x3 @cpu(0)>

If segment ids can remain fixed, you can make it a hybrid block too. Again, here is a quick and dirty implementation:

class SegmentSum(mx.gluon.HybridBlock):
    
    def __init__(self, sids, **kwargs):
        super(SegmentSum, self).__init__(**kwargs)
        self.sids = sids
    
    def hybrid_forward(self, F, x):
        
        ln = max(self.sids)
        
        sumargs = {}
        for i in range(ln+1):
            sumargs[i] = []

        for i in range(len(self.sids)):
            sumargs[self.sids[i]].append(F.slice_axis(x, axis=0, begin=i, end=i+1))

        sums = []

        for i in range(ln+1):
            sums.append(F.ElementWiseSum(*sumargs[i]))

        ans = F.stack(*sums)
        
        return ans

You can call it like:

sids = [0,0,1]

ss = SegmentSum(sids)
ss.hybridize()

ss.collect_params().initialize()

x = nd.array([[1,2,3],[-1,-2,-3],[5,6,7]])

y = ss(x)
print(y)

Output:

[[[ 0.  0.  0.]]

 [[ 5.  6.  7.]]]
<NDArray 2x1x3 @cpu(0)>

Since we are using the ElementWiseSum operator, it should be more efficient that calling add multiple times. Hope that helps.

1 Like

The basis is the same yes, though for cases with a segment of rank 2 (let’s say [3,12]), you would still need to do things sequentially in order to iterate over the first dimensions and the values present in the second, the use of elementwisesum is interesting, I’ll try to think of something inspired by your code and post it here.

Thanks for the help!