Max operation for specific elements in 4 dimensional array in parallel way

hohohihi: <p>I am coding pytorch. Between the torch inference code, I add some peripheral code for my own interest. This code works fine, but it is too slow. The reason might be for-iteration. So, i need parallel and fast way of doing this.</p>
<p>It is okay to do this in tensor, numpy, or just python array.</p>
<p>I made a function named ‘selective_max’ to find maximum value in arrays. But the problem is that I don’t want a maximum among the whole arrays, but among specific candidates which is designated by mask array. Let me show the gist of this function(below shows the code itself)</p>
<h1>input</h1>
<p>x [batch_size , dim, num_points, k] : x is a original input, but this becomes [batch_size, num_points, dim, k] by ‘x.permute(0,2,1,3)’.</p>
<p>‘batch_size’ is a well-known definition in the deep learning society. In every mini batch, there is many points. And a single point is represented by ‘dim’ length feature. Each feature element, there is ‘k’ potential candidates which is target of ‘max function’ later.</p>
<p>mask [batch_size, num_points, k] : This array is similar with ‘x’ without ‘dim’. Its element is either ‘0’ or ‘1’. So, I use this as a mask signal, like do max operation only on ‘1’ masked value.</p>
<p>please see the code below with this explanation. I use 3 for-iteration. Lets say we target a specific batch and a specific point. For a specific batch and a specific point, ‘x’ has [dim, k] array. And mask has [k] array which consists of either ‘0’ or ‘1’. So, I extract the non-zero index from [k] array and use this for extracting specific elements in ‘x’ dim by dim(‘for k in range(dim)’).</p>
<h1>toy example</h1>
<p>Let’s say we are in the second for-iteration. So, we now have [dim, k] for ‘x’ and [k] for ‘mask’. For this toy example, i presume k=3 and dim=4. x = [[3,2,1],[5,6,4],[9,8,7],[12,11,10], k=[0,1,1]. So, output would be [2,6,8,11], not [3, 6, 9, 12].</p>
<h1>previous try</h1>
<p>I try { mask.repeat(0,0,1,0) *(element-wise mul) x } and do the max operation. But, ‘0’ might the max value, because the x might have minus values in all array. So, this would result in wrong operation.</p>
<p>Thank you in advance.</p>
<pre><code class="lang-auto">def selective_max2(x, mask): # x : [batch_size , dim, num_points, k] , mask : [batch_size, num_points, k]
batch_size = x.size(0)
dim = x.size(1)
num_points = x.size(2)
k = x.size(3)
device = torch.device('cuda')

x = x.permute(0,2,1,3) # : [batch, num_points, dim, k]
#print('permuted x dimension : ',x.size())

x = x.detach().cpu().numpy()
mask = mask.cpu().numpy()
output = np.zeros((batch_size,num_points,dim))

for i in range(batch_size):
 for j in range(num_points):
  query=np.nonzero(mask[i][j]) # among mask entries, we get the index of nonzero values.
  for k in range(dim): # for different k values, we get the max value.
   # query is index of nonzero values. so, using query, we can get the values that we want.
   output[i][j][k] = np.max(x[i][j][k][query])

output = torch.from_numpy(output).float().to(device=device)
output = output.permute(0,2,1).contiguous()
return output```</code></pre>

KFrank: <p>Hihi Hoho!</p>
<aside class="quote no-group" data-username="hohohihi" data-post="1" data-topic="107487" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/bbce88/48.png" class="avatar"> hohohihi:</div>
<blockquote>
<h1>previous try</h1>
<p>I try { mask.repeat(0,0,1,0) *(element-wise mul) x } and do the max operation. But, ‘0’ might the max value, because the x might have minus values in all array. So, this would result in wrong operation.</p>
</blockquote>
</aside>
<p>Rather than use multiplication to set your masked elements to zero<br>
(which, as you note, will be incorrect), use subtraction to set your<br>
masked values to a large negative value.</p>
<p>Please see this “masked argmax” post:</p>
<aside class="quote quote-modified" data-post="4" data-topic="105341">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/masked-argmax-in-pytorch/105341/4">Masked argmax in Pytorch?</a> 
  </div>
  <blockquote>
    Hello Dragon! 

This won’t work if tensor t is negative (or, more precisely, if its 
largest unmasked element is negative). 
I would do this: 
large = torch.finfo (t.dtype).max   # assumes t is a kind of float
# assume msk has zeros where elements t should be masked out
# and ones where they should be kept
(t - large * (1 - msk) - large * (1 - msk)).argmax()

Best. 
K. Frank
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

hohohihi: <p>Hello, Frank. Thanks for reply. I see your point and it makes sense. let you know the result.</p>