Find elements surrounding an element in 2D tensor

username: Tankwell
cooked: <p>Hello</p>
<p>Suppose I have the following 2D tensor</p>
<pre><code class="lang-auto"> [[20, 29,  9,  7, 31, 35]],

        [[ 9, 23, 30,  5,  2, 31]],

        [[ 6, 27, 19,  5, 30,  7]],

        [[29,  4, 36,  2, 20, 27]],

        [[32,  2, 38, 10, 35,  1]],

        [[27, 40, 13, 10, 20,  3]],

        [[ 6, 20, 30,  7, 40, 39]]
</code></pre>
<p>I would like to create a new tensor, which will have the shape (…original_shape…,4)<br>
And will contain the 4 surrounding neighbors of each element, with a special index to the border.<br>
The directions will be “hard coded” (top,right,down,left)<br>
For example,</p>
<pre><code class="lang-auto">new_tensor[0,0,:] = [[[
BORDER,29,BORDER,BORDER
]]]
</code></pre>
<p>I thought on 2 options:</p>
<ul>
<li>
<p>Create <code>surround</code> function that will return the 4 surrounding elements (including BORDER in the right places) and run it on all tensor elements, and then build a new tensor from the results. But, as you may think, this is not the most efficient thing to do…</p>
</li>
<li>
<p>Use other tools, scatter / gather / others</p>
</li>
</ul>
<p>Do you have any recommendation?</p>
<p>Thanks<br>
Tankwell</p>

username: KFrank
cooked: <p>Hi Tankwell!</p>
<aside class="quote no-group" data-username="Tankwell" data-post="1" data-topic="162093" data-full="true">
<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/tankwell/48/53046_2.png" class="avatar"> Tankwell:</div>
<blockquote>
<p>I would like to create a new tensor, which will have the shape (…original_shape…,4)<br>
And will contain the 4 surrounding neighbors of each element, with a special index to the border.</p>
</blockquote>
</aside>
<p>Pad your input tensor with your special “border value” and then run <code>conv2d()</code><br>
with the appropriately constructed kernel:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; t = torch.tensor ([[20, 29,  9,  7, 31, 35],
...                    [ 9, 23, 30,  5,  2, 31],
...                    [ 6, 27, 19,  5, 30,  7],
...                    [29,  4, 36,  2, 20, 27],
...                    [32,  2, 38, 10, 35,  1],
...                    [27, 40, 13, 10, 20,  3],
...                    [ 6, 20, 30,  7, 40, 39]],
...                   dtype = torch.float)
&gt;&gt;&gt;
&gt;&gt;&gt; border_value = -99.0
&gt;&gt;&gt;
&gt;&gt;&gt; # build convolution kernel
&gt;&gt;&gt; kernel = torch.zeros (4, 1, 3, 3)
&gt;&gt;&gt; kernel[0, 0, 0, 1] = 1.0   # top
&gt;&gt;&gt; kernel[1, 0, 1, 2] = 1.0   # right
&gt;&gt;&gt; kernel[2, 0, 2, 1] = 1.0   # down
&gt;&gt;&gt; kernel[3, 0, 1, 0] = 1.0   # left
&gt;&gt;&gt;
&gt;&gt;&gt; tpad = torch.nn.functional.pad (t, (1, 1, 1, 1), mode = 'constant', value = border_value)
&gt;&gt;&gt; tsur = torch.nn.functional.conv2d (tpad.unsqueeze (0), kernel).permute (1, 2, 0)
&gt;&gt;&gt;
&gt;&gt;&gt; tsur.shape
torch.Size([7, 6, 4])
&gt;&gt;&gt; tsur[0, 0, :]
tensor([-99.,  29.,   9., -99.])
&gt;&gt;&gt; tsur
tensor([[[-99.,  29.,   9., -99.],
         [-99.,   9.,  23.,  20.],
         [-99.,   7.,  30.,  29.],
         [-99.,  31.,   5.,   9.],
         [-99.,  35.,   2.,   7.],
         [-99., -99.,  31.,  31.]],

        [[ 20.,  23.,   6., -99.],
         [ 29.,  30.,  27.,   9.],
         [  9.,   5.,  19.,  23.],
         [  7.,   2.,   5.,  30.],
         [ 31.,  31.,  30.,   5.],
         [ 35., -99.,   7.,   2.]],

        [[  9.,  27.,  29., -99.],
         [ 23.,  19.,   4.,   6.],
         [ 30.,   5.,  36.,  27.],
         [  5.,  30.,   2.,  19.],
         [  2.,   7.,  20.,   5.],
         [ 31., -99.,  27.,  30.]],

        [[  6.,   4.,  32., -99.],
         [ 27.,  36.,   2.,  29.],
         [ 19.,   2.,  38.,   4.],
         [  5.,  20.,  10.,  36.],
         [ 30.,  27.,  35.,   2.],
         [  7., -99.,   1.,  20.]],

        [[ 29.,   2.,  27., -99.],
         [  4.,  38.,  40.,  32.],
         [ 36.,  10.,  13.,   2.],
         [  2.,  35.,  10.,  38.],
         [ 20.,   1.,  20.,  10.],
         [ 27., -99.,   3.,  35.]],

        [[ 32.,  40.,   6., -99.],
         [  2.,  13.,  20.,  27.],
         [ 38.,  10.,  30.,  40.],
         [ 10.,  20.,   7.,  13.],
         [ 35.,   3.,  40.,  10.],
         [  1., -99.,  39.,  20.]],

        [[ 27.,  20., -99., -99.],
         [ 40.,  30., -99.,   6.],
         [ 13.,   7., -99.,  20.],
         [ 10.,  40., -99.,  30.],
         [ 20.,  39., -99.,   7.],
         [  3., -99., -99.,  40.]]])
</code></pre>
<p>(If you can live with <code>0.0</code> as your special border value, you can avoid<br>
the cost in time and memory of constructing <code>tpad</code> by using <code>Conv2d</code> with<br>
<code>padding_mode = 'zeros'</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Tankwell
cooked: <p>Hey,<br>
Tha’ts relly helped!</p>
<p>Actually I have to select some of the centers and not all of them.<br>
I have 2 ways to do that:</p>
<ol>
<li>Add them to the surround tensor above by 5 nonzero elements [add (1,1)] in the kernel, and looking for that value later</li>
<li>Use the original 2D tensor to find [for each sample] the i,j location of the <code>argument</code> passed (the value we are looking for). After, searching in that location on <code>surr</code>
</li>
</ol>
<p>The thing is that I have to locate a value in a 2D tensor on each sample in the batch…seems to be not so efficient</p>