Feature request: segment_argmax

[RanBenIzhak]

Currently trying to return argmax from a tensor for selected slices (segments)
slices do not have the same length, so reshaping can't be used.
was looking for a function similar to tf.segment_max, only with indices as the return value.

for e.g
a = [1, 2, 3, 4, 5, 6]
seg = [0, 0, 0, 1, 1, 2]
tf.segment_argmax return value will be
[2, 4, 5]

[tatatodd]

Contributions are welcome!

[srhrshr]

I can try to implement this.

Per the example by OP, the segments are [1,2,3], [4,5], [6] and the segment_argmax function returns the indices of max element in each segment, but with a reference to the unsegmented list.

Is my understanding correct ?

[RanBenIzhak]

yes,
indices will be with reference to the unsegmented list.

my intention was to use these indices to slice the original tensor for all the max tensors

[michaelzqj]

I will go to work on this.

[damro]

Any update on this? I would also extend the feature request to support unsorted_segment_argmax, segment_argmax, unsorted_segment_argmin, and segment_argmin.

[patnaa2]

Would also like to throw my vote on this as a feature ; this would be super useful.

[github-actions]

This issue is stale because it has been open for 180 days with no activity. It will be closed if no further activity occurs. Thank you.

[damro]

How is this progressing?