How do do matrix multiplication (matmal) along certain axis?

zeyuyun1: <p>I am relative new to pytorch. After doing a pretty exhaustive search online, I still couldn’t obtain the operation I want.</p>
<p>My question is How do do matrix multiplication (matmal) along certain axis?</p>
<p>For example, if I want to multiply a vector by a matrix, that would just be the following:<br>
a = torch.rand(3,5)<br>
b = torch.rand(3)<br>
torch.matmul(b,a)</p>
<p>One can interpret this as each element in b scale each row of a, and summing those scaled row together.</p>
<p>What if we have the dimension of a and b as following:<br>
a = torch.rand(3,5,10)<br>
b = torch.rand(3,10)<br>
and we want to do matrix multiplication along the first axis, basically, here’s what I want to do in a for loop form:</p>
<p>product = []<br>
for i in range(10):<br>
a_i = a[:,:,i]<br>
b_i = b[:,i]<br>
a_i_mul_b_i = torch.matmul(b_i,a_i)<br>
product.append(a_i_mul_b_i)</p>
<p>Although product is not a tensor but a list of tensor, but if we ignore the datatype, this is the end result I want.</p>

KFrank: <p>Hello Zeyuyun!</p>
<aside class="quote no-group quote-modified" data-username="zeyuyun1" data-post="1" data-topic="100867" 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/zeyuyun1/48/32747_2.png" class="avatar"> zeyuyun1:</div>
<blockquote>
<p>How do do matrix multiplication (matmal) along certain axis?<br>
…<br>
here’s what I want to do in a for loop form:</p>
<p>product = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
for i in range(10):<br>
a_i = a[:,:,i]<br>
b_i = b[:,i]<br>
a_i_mul_b_i = torch.matmul(b_i,a_i)<br>
product.append(a_i_mul_b_i)</p>
</blockquote>
</aside>
<p>The general-purpose tool for taking a product of (<em>contracting</em>) multiple<br>
tensors along various axes is <a href="https://pytorch.org/docs/1.7.0/generated/torch.einsum.html#torch.einsum" rel="noopener nofollow ugc">torch.einsum()</a> (named after “Einstein<br>
summation”).</p>
<p>(You can also fiddle with the dimensions to get them to line up as<br>
needed and use <code>matmul()</code> or <code>bmm()</code>.)</p>
<p>Here is a script that compares your loop code to <code>einsum()</code> (as well<br>
as to <code>bmm()</code> and <code>matmul()</code>):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__
torch.random.manual_seed (2020)

a = torch.rand(3,5,10)
b = torch.rand(3,10)

product = []
for i in range(10):
    a_i = a[:,:,i]
    b_i = b[:,i]
    a_i_mul_b_i = torch.matmul(b_i,a_i)
    product.append(a_i_mul_b_i)

# make product a tensor
product = torch.stack (product)

einsum_prod = torch.einsum ('ijk, ik -&gt; kj', a, b)
print ('check einsum_prod ...\n', torch.eq (einsum_prod, product).all())

matmul_prod = torch.matmul (b.unsqueeze (1).transpose (0, 2), a.permute (2, 0, 1)).squeeze()
print ('check matmul_prod ...\n', torch.eq (matmul_prod, product).all())

bmm_prod = torch.bmm (b.unsqueeze (1).transpose (0, 2), a.permute (2, 0, 1)).squeeze()
print ('check bmm_prod ...\n', torch.eq (bmm_prod, product).all())

print ('product = ...\n', product)
</code></pre>
<p>Here is its output</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.6.0'
&gt;&gt;&gt; torch.random.manual_seed (2020)
&lt;torch._C.Generator object at 0x7f8997a636f0&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.rand(3,5,10)
&gt;&gt;&gt; b = torch.rand(3,10)
&gt;&gt;&gt;
&gt;&gt;&gt; product = []
&gt;&gt;&gt; for i in range(10):
...     a_i = a[:,:,i]
...     b_i = b[:,i]
...     a_i_mul_b_i = torch.matmul(b_i,a_i)
...     product.append(a_i_mul_b_i)
...
&gt;&gt;&gt; # make product a tensor
&gt;&gt;&gt; product = torch.stack (product)
&gt;&gt;&gt;
&gt;&gt;&gt; einsum_prod = torch.einsum ('ijk, ik -&gt; kj', a, b)
&gt;&gt;&gt; print ('check einsum_prod ...\n', torch.eq (einsum_prod, product).all())
check einsum_prod ...
 tensor(True)
&gt;&gt;&gt;
&gt;&gt;&gt; matmul_prod = torch.matmul (b.unsqueeze (1).transpose (0, 2), a.permute (2, 0, 1)).squeeze()
&gt;&gt;&gt; print ('check matmul_prod ...\n', torch.eq (matmul_prod, product).all())
check matmul_prod ...
 tensor(True)
&gt;&gt;&gt;
&gt;&gt;&gt; bmm_prod = torch.bmm (b.unsqueeze (1).transpose (0, 2), a.permute (2, 0, 1)).squeeze()
&gt;&gt;&gt; print ('check bmm_prod ...\n', torch.eq (bmm_prod, product).all())
check bmm_prod ...
 tensor(True)
&gt;&gt;&gt;
&gt;&gt;&gt; print ('product = ...\n', product)
product = ...
 tensor([[0.8785, 0.5736, 0.3109, 0.5423, 0.6269],
        [0.7494, 1.4107, 0.9018, 1.1483, 1.2441],
        [0.9368, 0.9044, 0.9225, 0.9634, 0.1672],
        [0.1231, 0.4003, 0.4123, 0.1015, 0.4601],
        [0.9209, 1.1959, 0.5452, 1.0301, 0.8842],
        [0.2887, 1.1165, 0.9888, 0.8110, 0.4526],
        [0.9565, 1.2474, 0.1234, 1.6425, 0.9914],
        [1.1996, 0.8205, 1.0448, 1.3298, 1.0197],
        [0.3437, 0.3698, 0.4044, 0.4140, 0.8048],
        [0.7642, 1.4165, 0.7622, 0.6675, 1.2127]])
&gt;&gt;&gt;
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

zeyuyun1: <p>Thank you so much!! That’s exactly what I want. I think the key I am missing is “unsqueeze”, but eigenstein summation is another really good interpretation!</p>