Segmented operators (segment_sum)

I’m looking for an equivalent operator as tf.segment_sum, I couldn’t find any in the docs. If it’s not implemented, then this is a feature request.

Similarly I need the equivalent of tf.unsorted_segment_sum. Such segment reduction ops were implemented in Caffe2, which seem to be getting merged into PyTorch: https://github.com/pytorch/pytorch/blob/master/caffe2/operators/segment_reduction_op.cc

Hopefully they’ll be available in PyTorch 1.0…

Well, I just tried my hands on the segment_sum. You can check if out here: https://gist.github.com/bbrighttaer/207dc03b178bbd0fef8d1c0c1390d4be
You can let me know about its correctness after your verification.
Thanks.

1 Like

I’d like to point to a pytorch_scatter library, that implemements this functionality: https://github.com/rusty1s/pytorch_scatter