PyTorch - how to find the max values at certain indices ? (PyTorch equivalent of tf.math.unsorted_segment_max)

username: cdahms
cooked: <p>I posted this question on StackOverflow <a href="https://stackoverflow.com/questions/76335620/pytorch-how-to-find-the-max-values-at-certain-indices-pytorch-equivalent-o" class="inline-onebox" rel="noopener nofollow ugc">PyTorch - how to find the max values at certain indices ?? (PyTorch equivalent of tf.math.unsorted_segment_max) - Stack Overflow</a> but have not gotten any suggestions so I’ll post again here.</p>
<p>In PyTorch, I need to efficiently find the max values at certain indices. TensorFlow has a built in function <code>tf.math.unsorted_segment_max</code> to achieve this, I’ll borrow an example from their documentation <a href="https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_max:" rel="noopener nofollow ugc">https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_max:</a></p>
<pre><code class="lang-auto">import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf

data = tf.constant([5, 1, 7, 2, 3, 4], dtype=tf.float32)
segment_idxs = tf.constant([0, 0, 1, 1, 0, 1], dtype=tf.int64)
num_segments = 2

max_at_idxs = tf.math.unsorted_segment_max(data, segment_idxs, num_segments)

print('\n' + 'max_at_idxs: ')
print(max_at_idxs)
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">max_at_idxs: 
tf.Tensor([5. 7.], shape=(2,), dtype=float32)
</code></pre>
<p>I realize that logic-wise I could use a loop like this:</p>
<pre><code class="lang-auto">import torch

data = torch.tensor([5, 1, 7, 2, 3, 4], dtype=torch.float32)
segment_idxs = torch.tensor([0, 0, 1, 1, 0, 1], dtype=torch.int64)
num_segments = 2

max_at_idxs = torch.full(size=(num_segments, ), fill_value=-9999.9999, dtype=torch.float32)
for i, segment_idx in enumerate(segment_idxs):
    max_at_idxs[segment_idx] = torch.max(data[i], max_at_idxs[segment_idx])
# end for

print('\n' + 'max_at_idxs: ')
print(max_at_idxs)
</code></pre>
<p>However in my actual application <code>data</code> and <code>segment_idxs</code> are ~200,000 items long and this has to be done often so a loop is way too slow.</p>
<p>Similar questions have been asked in the PyTorch forum before, noteably: <a href="https://discuss.pytorch.org/t/how-to-perform-segment-max/33374" class="inline-onebox">How to perform segment max?</a><br>
<a href="https://discuss.pytorch.org/t/is-there-a-pytorch-implementation-of-tensor-flow-segment-softmax/77433/4" class="inline-onebox">Is there a pytorch implementation of tensor flow segment softmax? - #4 by tom</a><br>
<a href="https://discuss.pytorch.org/t/pytorch-equivalent-to-tf-unsorted-segment-sum/25275" class="inline-onebox">Pytorch equivalent to `tf.unsorted_segment_sum`</a><br>
but none of these resolves my concern, so I figured I’d ask here.</p>
<p>I keep figuring there has to be a way to do this with fancy indexing, possibly in combination with <code>torch.where</code> or some other built-in function, but I can’t seem to work it out.</p>
<p>Also, I keep looking at the <code>torch.maximum</code> function, trying to figure a way to use the 2nd parameter as indices rather than values, but I can’t find a way to work this out either.</p>
<p>Any suggestions?</p>

username: cdahms
cooked: <p>DerekG posted a great response <a href="https://stackoverflow.com/a/76335855/4835204" class="inline-onebox" rel="noopener nofollow ugc">PyTorch - how to find the max values at certain indices ?? (PyTorch equivalent of tf.math.unsorted_segment_max) - Stack Overflow</a>, I’ll repeat the answer here using my own variable namings in case the SO link ever goes bad:</p>
<pre><code class="lang-auto">import torch

xs = torch.tensor([8, 2, 6, 1, 3, 9, 4, 5, 7], dtype=torch.float32)
group_idxs = torch.tensor([0, 0, 1, 1, 2, 2, 0, 1, 2], dtype=torch.int64)
num_groups = 3
xs_len = len(xs)

# Slow way, un-comment if desired
# max_by_group = torch.full(size=(num_groups, ), fill_value=-9999.9999, dtype=torch.float32)
# for i, x in enumerate(xs):
#     max_by_group[group_idxs[i]] = max(x, max_by_group[group_idxs[i]])
# # end for

xs_exp = xs.unsqueeze(0).expand(num_groups, xs_len)

print('\n' + 'xs_exp: ')
print(xs_exp)

group_idxs_exp = group_idxs.unsqueeze(0).expand(num_groups, xs_len)

print('\n' + 'group_idxs_exp: ')
print(group_idxs_exp)

row_idxs = torch.arange(start=0, end=num_groups).unsqueeze(1).expand(num_groups, xs_len)

print('\n' + 'row_idxs: ')
print(row_idxs)

zero_or_neg_inf = torch.where(group_idxs_exp == row_idxs, 0, -torch.inf)

print('\n' + 'zero_or_neg_inf: ')
print(zero_or_neg_inf)

xs_exp_masked = xs_exp + zero_or_neg_inf

print('\n' + 'xs_exp_masked: ')
print(xs_exp_masked)

max_by_group, _ = torch.max(xs_exp_masked, dim=1)

print('\n' + 'max_by_group: ')
print(max_by_group)

print('\n')
</code></pre>