Matrix operations with PyTorch

Hdk: <p>Hello.</p>
<p>I was previously using several for loops for below use case, but it’s very very slow because I have bigger shapes of tensors. I saw the documentation of <code>torch.where</code>, but the pattern of tensors is very irregular, so I am stuck.<br>
I would be thankful if someone can provide a dummy example. <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>
<p>Use case:<br>
I have two 2D tensors say A and B, both of same shape.<br>
Tensor-A : say <code>torch.randn(5,5)</code> is</p>
<p>[ 0.1923,  1.6150, -0.4331,  0.7061,  1.7127],<br>
[-0.4912,  0.5317, -1.0820, -0.2575,  0.3446],<br>
[ 0.3956,  0.7645, -0.7015,  0.0574,  0.6930],<br>
[ 0.4492,  0.0215, -1.1855, -1.1453,  0.3912],<br>
[-0.5674, -0.8794, -0.1316,  0.4391, -1.0830]</p>
<p>Tensor-B consists of <strong>only three unique integer values</strong> (say 1,2,3). <strong>Each row</strong> consists of some 1s, 2s and 3s but in <strong>different proportion</strong>. Say, Tensor-B by<code> torch.randint(1,4,(5,5))</code> is<br>
[3, 1, 1, 1, 2],<br>
[3, 1, 2, 3, 3],<br>
[3, 2, 2, 1, 1],<br>
[2, 2, 3, 3, 1],<br>
[2, 3, 3, 1, 1]<br>
<code>##########</code></p>
<p>What I want to do is,</p>
<blockquote>
<blockquote>
<p><strong>For each row</strong>,</p>
<ol>
<li>Sum all the values of Tensor-A corresponding to the indices of all 1s of Tensor-B.</li>
<li>Divide each value of Tensor-A corresponding to the indices of all 2s of Tensor-B by the sum from previous step.</li>
<li>Take exponent after division for those values corresponding to all 2s and then sum again the values corresponding to all 2s for each row.</li>
</ol>
</blockquote>
</blockquote>
<p>i.e for above tensors if I take 3rd rows of A and B:</p>
<blockquote>
<blockquote>
<p>0.0574+0.6930 = 0.7504<br>
(0.7645/0.7504)= 1.019, (-0.7015/0.7504)= -0.9348<br>
exp(1.019)+exp(-0.9348) = 3.163</p>
</blockquote>
</blockquote>
<blockquote>
<blockquote>
<p>Similarly after doing the same for each row: Row1’s value+ Row2’s+ 3.163+…</p>
</blockquote>
</blockquote>

Emre_Yalcinoglu: <p>Hi!</p>
<p>I don’t know any compact way to achieve this, however, using the function <code>torch.mul</code> might be useful for this task. Here’s a possible solution:</p>
<pre><code class="lang-auto">A = torch.randn(5,5)
B = torch.randint(1,4,(5,5)).float()
Z = torch.zeros(5,5).float()

# I_n is a matrix where I_n(i,j) = 1 iff B(i,j)=n and for else I_n(i,j)=0
I1 = torch.where(B==1, B, Z)
I2 = torch.where(B==2, torch.ones(5,5), Z)
I3 = torch.where(B==3, torch.ones(5,5), Z)

# observe that torch.mul is different than torch.mm
step_1 = torch.sum(torch.mul(I1, A),dim=1)

# reshape for transposing
pre_step_2 = step_1.reshape(-1,1).repeat(1,5)
step_2 = torch.mul(torch.mul(I2, A), pre_step_2)

# here we also sum the components which are not corresponding to 2,
# which are 0, so they will become 1 under exp 
step_3 = torch.sum(torch.exp(step_2))

result = step_3-torch.sum(I1+I3)
</code></pre>

Hdk: <p>Thanks <a class="mention" href="/u/emre_yalcinoglu">@Emre_Yalcinoglu</a> for the reply.</p>
<p>I have one more confusion.<br>
Say we have a Tensor-B of shape 12x12. (Which we can imagine, is divided in small matrices of 4x4 both row and columnwise.) And initially B contains only 1s and 2s.</p>
<p>Now, I want to replace all the entries of main diagonal+(5th ,6th ,7th ,9th ,10th ,11th) upper and lower diagonals, by 3s. I don’t know what would be a faster way.</p>

Emre_Yalcinoglu: <p>See <a href="https://pytorch.org/docs/stable/generated/torch.diag_embed.html" rel="noopener nofollow ugc">torch.diag_embed</a> and <a href="https://pytorch.org/docs/stable/generated/torch.diagonal.html" rel="noopener nofollow ugc">torch.diagonal</a>. You can simulate replacement by <code>torch.where</code>.<br>
I would suggest writing your own short <a href="https://docs.python.org/3/reference/expressions.html#lambda" rel="noopener nofollow ugc">lambda expressions</a> to handle such tasks.</p>