I have a tensor of shape
(batch_size, num_steps, hidden_size). For each row vector in each matrix
(num_steps, hidden_size) I want to compute the element-wise product with every other row vector, yielding a tensor of shape
(batch_size, num_steps, num_steps, hidden_size).
Is there a built-in way to do this that avoids for loops?
You can use broadcasting. For example:
y = x[:, :, None, :] * x[:, None, :, :]
There also is torch.einsum for fancy stuff (
torch.einsum("iaj,ibj->iabj", [x, x]), works in PyTorch self-compiled master), but in your case the above seems to be what looks most straightforward and transparent to me (the trailing “:” indices are for my taste, not technically necessary, alternatives include inserting the dimensions with unsqueeze instead of using indexing).
That’s brilliant. Quite a saving from my 10 lines of for loop rubbish. It has also expanded my world in terms of what can be achieved with broadcasting. I’ve played with it for a bit now and am thoroughly satisfied with the solution.
Thanks very much for your kind help!!!