Pytorch Implementation of Tensorflow code

Ajinkya_Ambatwar: <p>Hi I am trying to move a TF code to PyTorch. But I am unable to replicate one functionality into Torch. I am pasting the two lines from tensorflow here</p>
<pre><code class="lang-auto">neighbourRagged = tf.RaggedTensor.from_row_splits(x, row_splits=rowSplits)
neighbourFeatures = tf.math.reduce_max(neighbourRagged, axis=1)
</code></pre>
<p>where x is a huge float array of size (257997,64) and row_splits is an int32 array of size (32769,).</p>
<p>I am trying to implement these two lines from tf to torch. I am aware that Torch does not support <code>RaggedTensor</code>  yet. But is there any work-around that to “specifically” implement these two lines?<br>
I have uploaded the exact <code>.npy</code> files(zipped) on drive which can be accessed from <a href="https://1drv.ms/u/s!At142eQK9YsQ8g_4AOKegtw--R11" rel="noopener nofollow ugc">here</a></p>
<p>The <code>x</code> and <code>rowSplits</code> array can be accessed from “files” folder after extracting zip file.</p>
<pre><code class="lang-auto">x = np.load("x.npy")
rowSplits = np.load("numNeighboursArray.npy")
</code></pre>
<p>If you are able to replicate the functionality please do let me know!<br>
Thank you in advance!!</p>
<p>Tagging <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <p>I’m not familiar with the TF code, but would these two operations refer to a <code>scatter_max</code> op?<br>
If so, you could take a look at <a href="https://discuss.pytorch.org/t/how-to-implement-efficiently-scatter-max-in-pytorch/42990">this post</a>.<br>
Or is the reduction axis in the max operation different from the split axis and you would still get “ragged tensors” with scalar values in one dimension?</p>

Ajinkya_Ambatwar: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> sir,</p>
<p>The first line splits the x array (n,64) along the n axis based on indices given in <code>rowSplits</code> and converts it into a raggedTensor of size <code>(32768,[None],64)</code> where <code>None</code> in shape indicates variable number<br>
eg. <code>raggedTensor [0].shape=[20,64]</code> ,<code>raggedTensor [1].shape=[19,64]</code> etc.</p>
<p>The 2nd line takes the Max along axis 1( the variable length axis) to get the output. In the example given here, <code>op[0]</code> is the element wise maximum across the 20 samples obtained. <code>op[1]</code> is for 19 samples etc.<br>
I hope I have made it clearer.<br>
Can I replicate this behaviour using <code>scatter_max</code> or any other operation?<br>
<a class="mention" href="/u/alband">@albanD</a> can you please suggest something on this?<br>
Thank you!</p>

Ajinkya_Ambatwar: <p>Is it not possible to write this code in Torch?<br>
Not having RaggedTensor is one of the few downfalls I am facing with Torch. If there is some workaround, someone please help me through that.</p>

J_Johnson: <p>Have you tried just padding your matrix with zeros to make them all the same size? I don’t see the rational/benefit to making a matrix variable on one axis for ML.</p>

Ajinkya_Ambatwar: <p><a class="mention" href="/u/j_johnson">@J_Johnson</a> that can be done. But here is the caveat-<br>
Eg.</p>
<pre><code class="lang-auto">x=torch.randn(10,3)
RowSplits=[0,2,5,7,9]
</code></pre>
<p>Which means the splits are <code>x[0:2], x[2:5], x[5:7], x[7:9]</code><br>
Now to make each of them of same size, I shall have to pad zeros in between these rows. Say I write new x into x_ like</p>
<pre><code class="lang-auto">   x_[0:2]=x[0:2]
    x_[2]=0(inserted padding)
    x_[3:6]=x[2:5](no padding needed here) 
</code></pre>
<p>etc. Now how can I do it efficiently without having to run a for loop(performance hit)<br>
Another point to note is I have to take Max in each such partition. Like in original x array I need element-wise Max from x[0] to x[1] i.e<br>
(From first slice)<br>
<code>op[0,0]=max{x[0,0],x[1,0]}</code><br>
<code>op[0,1]=max{x[0,1],x[1,1]}</code><br>
<code>op[0,2]=max{x[0,2],x[1,2]}</code></p>
<p>(From second slice)<br>
<code>op[1,0]=max{x[2,0],x[3,0],x[4,0]}</code><br>
<code>op[1,1]=max{x[2,1],x[3,1],x[4,1]}</code><br>
<code>op[1,2]=max{x[2,2],x[3,2],x[4,2]}</code><br>
And so on</p>
<p>Can this be done efficiently?</p>

Ajinkya_Ambatwar: <p>Hi I found one workaround based on <a class="mention" href="/u/j_johnson">@J_Johnson</a>’s suggestion!<br>
I am showing the demo numpy implementation for the same.</p>
<pre><code class="lang-auto">x = np.random.randint(low=10,size=(10,3))
# array([[6, 0, 3],
       [0, 7, 0],
       [4, 6, 8],
       [1, 6, 7],
       [2, 6, 1],
       [8, 5, 5],
       [9, 8, 3],
       [4, 9, 1],
       [5, 8, 6],
       [6, 6, 3]])

splits = np.array([0,2,3,6,7,10])
splits_ = np.subtract(splits[1:],splits[:-1])
# splits_ = [2,1,3,1,3]
l = len(splits_) # l =5
m = max(splits_) # m = 3
neighbourRagged = np.zeros([l,m,3]) 
# this will be my output array where I can copy
# the input with zero padding
</code></pre>
<p>I set <code>output</code>'s second axis length to the maximum split value( in this case 3). This is the maximum value of the second axis after zero-padding</p>
<p>Now for each block(among <code>l</code> blocks) each of size <code>(m,3)</code>, I fill up the first <code>i</code> rows (obtained from each split value)   with the <code>i==j-prev</code> rows from input array and rest rows stay at zero.</p>
<pre><code class="lang-auto">prev = 0
idxs = range(l)
for idx,i,j in zip(idxs,splits_, splits[1:]):
...     neighbourRagged[idx, :i,:]=x[prev:j,:]
...     prev=j

#neighbourRagged
array([[[6., 0., 3.],
        [0., 7., 0.],
        [0., 0., 0.]],

       [[4., 6., 8.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[1., 6., 7.],
        [2., 6., 1.],
        [8., 5., 5.]],

       [[9., 8., 3.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[4., 9., 1.],
        [5., 8., 6.],
        [6., 6., 3.]]])
</code></pre>
<p>Now to obtain the max output, I can take the maximum along <code>axis=1</code></p>
<pre><code class="lang-auto">neighbourFeatures = neighbourRagged.max(axis=1)
#neighbourFeatures
array([[6., 7., 3.],
       [4., 6., 8.],
       [8., 6., 7.],
       [9., 8., 3.],
       [6., 9., 6.]])
</code></pre>
<p>This is the output I want.<br>
I was wondering if we can implement this in PyTorch using some efficient function.(without using loops etc.)</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a></p>

J_Johnson: <p>Maybe can try:</p>
<pre><code class="lang-auto">x = torch.randint(0, 10, (10,3))

splits = torch.tensor([0, 2, 3,6,7,10])

x_split = [x[splits[i]:splits[i+1]] for i in range(len(splits)-1)]
</code></pre>
<p>Or something to that effect, depending on the specifics.</p>

ptrblck: <p>I haven’t checked it, but have you had a chance to try e.g. <a href="https://github.com/rusty1s/pytorch_scatter">PyTorch Scatter</a> from the previously linked topic?<br>
If so, did a scatter operation work for you or were you hitting any issues with it?</p>

Ajinkya_Ambatwar: <aside class="quote no-group quote-modified" data-username="ptrblck" data-post="7" data-topic="114668">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If so, did a scatter operation work for you or were you hitting any issues with it?</p>
</blockquote>
</aside>
<p>Thank you so much <a class="mention" href="/u/ptrblck">@ptrblck</a> !<br>
The <code>segment_csr</code> function from <code>torch_scatter</code> is just exactly what I wanted.</p>
<pre><code class="lang-auto">from torch_tensor import segment_csr
x = torch.tensor([[6, 0, 3],
       [0, 7, 0],
       [4, 6, 8],
       [1, 6, 7],
       [2, 6, 1],
       [8, 5, 5],
       [9, 8, 3],
       [4, 9, 1],
       [5, 8, 6],
       [6, 6, 3]])
splits = torch.tensor([0,2,3,6,7,10])
segment_csr(x, splits,reduce='max')
# tensor([[6, 7, 3],
        [4, 6, 8],
        [8, 6, 7],
        [9, 8, 3],
        [6, 9, 6]])
# Exact answer
Thank you once again @ptrblck 
Thank you to @J_Johnson  as well</code></pre>