Batch hadamard product of every vector in a matrix with every other

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?

1 Like

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).

Best regards

Thomas

3 Likes

Dear Thomas,

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!!!

Best regards,

Tim.

1 Like