Most efficient way to perform multiplication of list of tensors?

hugo_thim: <p>Hello everyone,</p>
<p>I am wondering what is the most efficient way to perform multiplication of tensors contained in a sublist within two lists.</p>
<p>Let us say that I have a list <code>A</code> containing sublists <code>a1,...ak</code> containing tensors of different sizes within each list, e.g. <code>a1 = [a11, a12, a13, a14]</code> where</p>
<pre><code class="lang-auto">a11.size() == torch.Size([128, 274])
a12.size() == torch.Size([1, 128])
a13.size() == torch.Size([256, 128])
a14.size() == torch.Size([1, 256])
</code></pre>
<p>but note that each list within <code>A</code> contains tensors that have the same sizes as those in <code>a1</code>.</p>
<p>Now consider a quite similar list, <code>B = [b1,..., bn]</code>, where each sublist <code>bi</code> contains the same number of tensors as those in the sublist of <code>A</code>. Moreover, the inner tensors have the same shapes as those in each sublist <code>aj</code>.</p>
<p>I am looking to “multiply” each sublist in <code>A</code> with each sublist in <code>B</code> such that I get a list <code>C</code> which contains <code>k</code> lists, each containing the matrix multplication of the sublists <code>aj</code> with each sublist in <code>B</code>. Formally,</p>
<pre><code class="lang-auto">C = [ [a1 * b1, a1 * b2, ..., a1 * bn], [a2 * b1, a2 * b2, ..., a2 * bn], ..., ]
</code></pre>
<p>where  <code>a1 * b1</code> denotes <code>[a11 * b11, a12 * b12, a13 * b13, a14 * b14]</code>.</p>
<p>However, I cannot perform a multiplication between list of tensors efficiently, for instance trying <code>a1 * b1</code> outputs the following error</p>
<pre><code class="lang-auto">TypeError: can't multiply sequence by non-int of type 'list'
</code></pre>
<p>Thus my guess is that the format in which I am storing my tensors is probably not optimal.</p>
<p>Is the best way to perform such computation the following</p>
<pre><code class="lang-auto">c =  [ [torch.mul(aik, bjk)) for aik, bjk in zip(ai, bj)] for ai, bj in zip(A, B)]
</code></pre>
<p>or is there a more efficient way to do so ? (e.g. storing tensors in other than a list).</p>

googlebot: <p>it is best to have multiplicands in a single memory block (buffer), but this is often impossible to arrange, as copying is not efficient, and “out” argument breaks gradients. if you mean matrix multiplication, you can’t “pack” tensors with varying shapes either.</p>
<p>apart from that, you can only reduce invocation overhead a bit (e.g. with JIT), but this will likely be insignificant, unless you’re dealing with hundreds of small tensors.</p>

hugo_thim: <p>I believe I have found the answer to my question. Since, tensors can only be concatenated when of identical size, I can simply concatenate each tensors of identical sizes in each of the lists in <code>A</code> and <code>B</code> to make 3D tensors.</p>
<p>For instance,</p>
<pre><code class="lang-auto">A = [torch.stack(x) for x in list(zip(*A))]
B = [torch.stack(x) for x in list(zip(*B))]
</code></pre>
<p>outputs two lists containing <code>4</code> 3D tensors each having respectively <code>k</code> and <code>n</code> as first dimension. To stay coherent with the post, the first element in <code>A</code> has now the following size</p>
<pre><code class="lang-auto">A[0].size() == torch.Size([k, 128, 274])
</code></pre>
<p>Then flattening each of the tensors in dimension 1 and 2 + transposition allows standard faster matrix multiplication.</p>
<pre><code class="lang-auto">C = [(torch.mm(torch.flatten(a, start_dim=1),
               torch.flatten(b, start_dim=1).transpose(0,1)), dim=1) 
      for a, b in zip(A, B)]
</code></pre>
<p>This allowed my code to run 6x faster.</p>