Chained broadcasting matmul without for loop

MrCrHaM: <p>I have multiple tensors mixed with <strong>batched matrices</strong> and <strong>broadcasted matrices</strong>.  Is there an efficient and elegant way to get the matrix product of them?</p>
<p>Following snippet is how it’s implemented currently which is ugly. BTW, for loop is better avoided due to performance reason.</p>
<pre><code class="lang-python">tensor1 = torch.randn(10, 3, 4)
tensor2 = torch.randn(4, 5)
tensor3 = torch.randn(10, 5, 6)
tensor4 = torch.randn(6, 7)
tensor5 = torch.randn(10, 7, 8)

result = torch.matmul(
    torch.matmul(
        torch.matmul(
            torch.matmul(
                tensor1,
                tensor2),
            tensor3),
        tensor4),
    tensor5,
)
print(result.size())  # torch.Size([10, 3, 8])
</code></pre>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/mrcrham">@MrCrHaM</a>,</p>
<p>You can get the desired output via using <code>torch.func.vmap</code> and <code>torch.linalg.multi_dot</code>, here’s a minimal reproducible example.</p>
<pre><code class="lang-auto">import torch
from torch.linalg import multi_dot
from torch.func import vmap

tensor1 = torch.randn(10, 3, 4)
tensor2 = torch.randn(4, 5)
tensor3 = torch.randn(10, 5, 6)
tensor4 = torch.randn(6, 7)
tensor5 = torch.randn(10, 7, 8)

def func(A,B,C,D,E):
  return multi_dot((A,B,C,D,E)) #notice the 2nd ()'s here.
  
out = vmap(func, in_dims=(0,None,0,None,0))(tensor1,tensor2,tensor3,tensor4,tensor5)
print(out.shape) #returns torch.Size([10, 3, 8])
</code></pre>

MrCrHaM: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>,</p>
<p>Thanks a lot for the great answer.</p>
<p>However, I’m currently using PyTorch 1.12.1. It seems that <code>torch.func</code> is not available, and upgrading to 2.0 would involve uprading CUDA and NVIDIA drive accordingly which is not an option for the cluster for now.</p>
<p>Any ideas to solve it without any PyTorch 2.0 specific features?</p>

AlphaBetaGamma96: <p>You can change <code>torch.func.vmap</code> to <code>functorch.vmap</code>, the <code>functorch</code> package should have been installed alongside <code>pytorch</code> when you initially installed it.</p>

ptrblck: <aside class="quote no-group" data-username="MrCrHaM" data-post="3" data-topic="187070">
<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/mrcrham/48/25149_2.png" class="avatar"> MrCrHaM:</div>
<blockquote>
<p>and upgrading to 2.0 would involve uprading CUDA and NVIDIA drive accordingly which is not an option for the cluster for now.</p>
</blockquote>
</aside>
<p>Why would this be the case? The PyTorch 2.0.0 binaries ship with CUDA 11.7 and 11.8 and all needed dependencies. You would this only need to install a proper NVIDIA driver, which support minor-version compatibility since CUDA 11. Were you using a CUDA 10.2 build?</p>

MrCrHaM: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="4" data-topic="187070" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>You can change <code>torch.func.vmap</code> to <code>functorch.vmap</code>, the <code>functorch</code> package should have been installed alongside <code>pytorch</code> when you initially installed it.</p>
</blockquote>
</aside>
<p>Sadly, it seems <code>functorch</code> package is not available for PyTorch 1.12.1 as well.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import functorch
ModuleNotFoundError: No module named 'functorch'
</code></pre>

MrCrHaM: <p>The cluster I’m running code on is using CUDA 10.2 and according to the administrator upgrading CUDA is not planned for now unfortunately. So I guess PyTorch 2.0 specific features won’t be available for me anytime soon. <img src="https://discuss.pytorch.org/images/emoji/apple/face_exhaling.png?v=12" title=":face_exhaling:" class="emoji" alt=":face_exhaling:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>Ah OK, that explains it. Major CUDA versions need a driver update (unless you are using the compat driver package on data center GPUs). Note that PyTorch does not support CUDA 10.2 anymore, so you would eventually need to update if you want to use any newer release.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/mrcrham">@MrCrHaM</a>,</p>
<p>I read the <code>functorch</code> docs, for your case, you can install <code>functorch</code> via <code>pip install functorch</code>, more info in the docs here: <a href="https://pytorch.org/functorch/stable/install.html" class="inline-onebox" rel="noopener nofollow ugc">Install functorch — functorch nightly documentation</a></p>
<p>You should be able to install with conda too if you need to, by just locally installing to pip inside your conda environment.</p>
<p>Also, you will get warned that you’re using a depracted package and to use <code>torch.func</code> instead, just make sure to suspress the warnings with the <code>warnings</code> package (should be installed already) and things should be fine</p>

MrCrHaM: <p>Thanks for the info. I’ll send these words to our cluster administrator. <img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=12" title=":joy:" class="emoji" alt=":joy:" loading="lazy" width="20" height="20"></p>

MrCrHaM: <p>Thanks for the further instruction. Module <code>functorch</code> works like a charm.</p>