How to compute outer product between vectors with broadcast and batch?

Wang-Yu-Qing: <p>Is there any operation that can achieve the following:</p>
<pre><code class="lang-auto">import torch

batch_size = 2
seq_len = 2
dim = 3
# batch of squences of embedding vecs:
x = torch.rand([batch_size, seq_len, dim])
# batch of target embedding vecs:
y = torch.rand([batch_size, dim])

# the computation I want to achieve:
print(torch.outer(x[0][0], y[0]))
print(torch.outer(x[0][1], y[0]))
print(torch.outer(x[1][0], y[1]))
print(torch.outer(x[1][1], y[1]))
print()

# what I've tried but failed:
#print(torch.einsum('bij, bj-&gt;bij', x, y))
# or
#print(torch.einsum('bij, bj-&gt;bijk', x, y))
</code></pre>

KFrank: <p>Hi Yu-Qing!</p>
<aside class="quote no-group" data-username="Wang-Yu-Qing" data-post="1" data-topic="137767" 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/wang-yu-qing/48/40608_2.png" class="avatar"> Wang-Yu-Qing:</div>
<blockquote>
<pre><code class="lang-auto"># the computation I want to achieve:
print(torch.outer(x[0][0], y[0]))
print(torch.outer(x[0][1], y[0]))
print(torch.outer(x[1][0], y[1]))
print(torch.outer(x[1][1], y[1]))
</code></pre>
</blockquote>
</aside>
<p><code>torch.einsum ('ijk, il -&gt; ijkl', x, y)</code> should do what you want:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; batch_size = 2
&gt;&gt;&gt; seq_len = 2
&gt;&gt;&gt; dim = 3
&gt;&gt;&gt; x = torch.rand([batch_size, seq_len, dim])
&gt;&gt;&gt; y = torch.rand([batch_size, dim])
&gt;&gt;&gt;
&gt;&gt;&gt; z = torch.empty (batch_size, seq_len, dim, dim)
&gt;&gt;&gt;
&gt;&gt;&gt; z[0, 0] = torch.outer(x[0][0], y[0])
&gt;&gt;&gt; z[0, 1] = torch.outer(x[0][1], y[0])
&gt;&gt;&gt; z[1, 0] = torch.outer(x[1][0], y[1])
&gt;&gt;&gt; z[1, 1] = torch.outer(x[1][1], y[1])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.einsum ('ijk, il -&gt; ijkl', x, y).equal (z)
True
&gt;&gt;&gt;
&gt;&gt;&gt; print(torch.outer(x[0][0], y[0]))
tensor([[0.3167, 0.2910, 0.2960],
        [0.7378, 0.6779, 0.6894],
        [0.6072, 0.5579, 0.5674]])
&gt;&gt;&gt; print(torch.outer(x[0][1], y[0]))
tensor([[0.3050, 0.2802, 0.2850],
        [0.0210, 0.0193, 0.0196],
        [0.2876, 0.2642, 0.2687]])
&gt;&gt;&gt; print(torch.outer(x[1][0], y[1]))
tensor([[0.4431, 0.4939, 0.6687],
        [0.4362, 0.4863, 0.6584],
        [0.2593, 0.2891, 0.3914]])
&gt;&gt;&gt; print(torch.outer(x[1][1], y[1]))
tensor([[0.3535, 0.3940, 0.5335],
        [0.3736, 0.4165, 0.5639],
        [0.1286, 0.1433, 0.1941]])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.einsum ('ijk, il -&gt; ijkl', x, y)
tensor([[[[0.3167, 0.2910, 0.2960],
          [0.7378, 0.6779, 0.6894],
          [0.6072, 0.5579, 0.5674]],

         [[0.3050, 0.2802, 0.2850],
          [0.0210, 0.0193, 0.0196],
          [0.2876, 0.2642, 0.2687]]],


        [[[0.4431, 0.4939, 0.6687],
          [0.4362, 0.4863, 0.6584],
          [0.2593, 0.2891, 0.3914]],

         [[0.3535, 0.3940, 0.5335],
          [0.3736, 0.4165, 0.5639],
          [0.1286, 0.1433, 0.1941]]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>