Applying a function on each individual element of a Tensor

username: osama-usuf
cooked: <p>I have a task where I perform an operation over individual elements of the network gradient.</p>
<p>For a fully-connected layer, I have an <code>m x n</code> gradient matrix <strong>G</strong>. I have an equivalent <code>m x n</code> update matrix <strong>P</strong>.</p>
<p>Elements of <strong>G</strong> and <strong>P</strong> map one-to-one i.e. that <strong>P_{ij}</strong> dictates how to update <strong>G_{ij}</strong></p>
<p>I have gotten this to work with a nested loop, but that was extremely slow. I managed to improve the speed with the multiprocessing <a href="https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.starmap" rel="noopener nofollow ugc">starmap</a> method, but I want to do this in a better, more direct way, as this is still slow and takes a long time.</p>
<p>Is there a PyTorch way to perform this operation in parallel? I understand that this is theoretically possible, but I have not been able to figure this out. I have also tried PyTorch’s <code>vmap</code> function with no success.</p>
<p>Here’s a re-creatable example of what I want to be able to achieve.</p>
<pre><code class="lang-auto">import torch
import numpy as np


def update_gradient(g, p):
    dummy = 1

    if p &gt; 1:
        dummy += 100
    else:
        dummy -= 9

    return dummy


m, n = 10, 15

gradient = torch.rand(m, n)
pulses = torch.rand(m, n)

new_list = []


# instead of this nested for loop - I want to map my update_gradient method in parallel to each element of the gradient and pulse matrix
for i in range(0, gradient.shape[0]):
    for j in range(0, gradient.shape[1]):
        new_list.append(update_gradient(gradient[i][j], pulses[i][j]))


new_list = torch.Tensor(new_list)
new_list = torch.reshape(new_list, gradient.shape)
</code></pre>

username: KFrank
cooked: <p>Hi Osama!</p>
<aside class="quote no-group" data-username="osama-usuf" data-post="1" data-topic="143621" 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/osama-usuf/48/46459_2.png" class="avatar"> osama-usuf:</div>
<blockquote>
<p>Is there a PyTorch way to perform this operation in parallel?</p>
</blockquote>
</aside>
<p>In general, if you want to apply a function element-wise to the elements<br>
of a pytorch tensor and that function is built up of “straightforward” pieces,<br>
it will usually be possible to rewrite that function in terms of pytorch tensor<br>
operations that work on the tensor as a whole (element-wise) without using<br>
loops.</p>
<p>Here’s an illustration that compares your example code with a loop-free<br>
version:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; def update_gradient(g, p):
...     dummy = 1
...
...     if p &gt; 1:
...         dummy += 100
...     else:
...         dummy -= 9
...
...     return dummy
...
&gt;&gt;&gt; m, n = 10, 15
&gt;&gt;&gt;
&gt;&gt;&gt; gradient = torch.rand(m, n)
&gt;&gt;&gt; pulses = torch.rand(m, n)
&gt;&gt;&gt;
&gt;&gt;&gt; new_list = []
&gt;&gt;&gt;
&gt;&gt;&gt; # instead of this nested for loop - I want to map my update_gradient method in parallel to each element of the gradient and pulse matrix
... for i in range(0, gradient.shape[0]):
...     for j in range(0, gradient.shape[1]):
...         new_list.append(update_gradient(gradient[i][j], pulses[i][j]))
...
&gt;&gt;&gt; new_list = torch.Tensor(new_list)
&gt;&gt;&gt; new_list = torch.reshape(new_list, gradient.shape)
&gt;&gt;&gt;
&gt;&gt;&gt; # loop-free, pytorch-tensor-function version
... def update_gradientB (g, p):
...     dummy = torch.ones_like (g)
...     dummy[p  &gt; 1] += 100
...     dummy[p &lt;= 1] -= 9
...     return dummy
...
&gt;&gt;&gt; resultB = update_gradientB (gradient, pulses)
&gt;&gt;&gt;
&gt;&gt;&gt; resultB
tensor([[-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.],
        [-8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8., -8.,
         -8.]])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (resultB, new_list)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: osama-usuf
cooked: <p>I understand your solution - but I guess my problem then is that the function I want to write in terms of PyTorch operations isn’t really re-writable as PyTorch operations (at least from what I have been able to work out).</p>
<p>I have asked a follow-up question <a href="https://discuss.pytorch.org/t/calculating-index-of-nearest-value-to-a-nested-tensor/143693">here</a>  particular to my problem - thanks for your input <a class="mention" href="/u/kfrank">@KFrank</a> !</p>