Broadcasting with Bilinear function

username: legoh
cooked: <p>Hi there, I would like some help with trying to do a vector outer product, but with a bilinear function. So for example, we have the following vectors:</p>
<pre><code class="lang-auto">a = torch.rand(1, 10, 128)  # |a| = (batch, num_vectors, dim)
b = torch.rand(1, 10, 128)  # |b| = (batch, num_vectors, dim)

# In a vector outer product, we could just do
c = torch.bmm(a, b.permute(0, 2, 1))  # |c| = (batch, num_vectors, num_vectors)

# I would like to do the same operation, but using a Bilinear layer
# However, if we use a Bilinear layer, we need to repeat vectors. Meaning it would look like
a = a.repeat_interleave(10, dim=1)
b = b.repeat(1, 10, 1)

bilinear_layer = torch.nn.Bilinear(128, 128, 1)
c = bilinear_layer(a, b) # |c| = (batch, num_vectors * num_vectors, 1)
c = c.view(batch, num_vectors, num_vectors)
</code></pre>
<p>The thing is, <code>repeat_interleave</code> and <code>repeat</code> will allocate memory. Is there a way to efficiently broadcast this operation and avoid that?</p>

username: KFrank
cooked: <p>Hi Legoh!</p>
<aside class="quote no-group quote-modified" data-username="legoh" data-post="1" data-topic="161868" 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/l/82dd89/48.png" class="avatar"> legoh:</div>
<blockquote>
<p>I would like some help with trying to do a vector outer product, but with a bilinear function.<br>
…<br>
The thing is, <code>repeat_interleave</code> and <code>repeat</code> will allocate memory. Is there a way to efficiently broadcast this operation and avoid that?</p>
</blockquote>
</aside>
<p>I do not believe that there is a way to tell (or trick) <code>Bilinear</code> to have<br>
it broadcast internally and I don’t believe that there is a way to reshape<br>
<code>a</code> and <code>b</code> appropriately without allocating memory.</p>
<p>I think the cleanest approach is to use <code>einsum()</code> with <code>Bilinear.weight</code><br>
to get, in effect, the broadcasting you want:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; batch = 2
&gt;&gt;&gt; num_vectors = 10
&gt;&gt;&gt; dim = 128
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.rand (batch, num_vectors, dim)  # |a| = (batch, num_vectors, dim)
&gt;&gt;&gt; b = torch.rand (batch, num_vectors, dim)  # |b| = (batch, num_vectors, dim)
&gt;&gt;&gt; aa = a.clone()
&gt;&gt;&gt; bb = b.clone()
&gt;&gt;&gt;
&gt;&gt;&gt; # I would like to do the same operation, but using a Bilinear layer
&gt;&gt;&gt; # However, if we use a Bilinear layer, we need to repeat vectors. Meaning it would look like
&gt;&gt;&gt; a = a.repeat_interleave (num_vectors, dim = 1)
&gt;&gt;&gt; b = b.repeat  (1, num_vectors, 1)
&gt;&gt;&gt;
&gt;&gt;&gt; bilinear_layer = torch.nn.Bilinear (dim, dim, 1)
&gt;&gt;&gt; c = bilinear_layer (a, b) # |c| = (batch, num_vectors * num_vectors, 1)
&gt;&gt;&gt; c = c.view (batch, num_vectors, num_vectors)
&gt;&gt;&gt;
&gt;&gt;&gt; cc = torch.einsum ('ijk, nkm, ilm -&gt; ijl', aa, bilinear_layer.weight, bb) + bilinear_layer.bias
&gt;&gt;&gt; cc.shape
torch.Size([2, 10, 10])
&gt;&gt;&gt; torch.allclose (cc, c)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: legoh
cooked: <p>Thanks Frank! I incorporated your solution and it works well! Thank you so much.</p>