Einsum for shapes of different sizes or ranks
I have two PyTorch tensors. One is rank three and the other is rank four. Is there a way to get it so that it produce the rank and shape of the first tensor? For instance in this cross-attention bit:
q = torch.linspace(1, 192, steps=192)
q = q.reshape(2, 4, 3, 8)
k = torch.linspace(2, 193, steps=192)
k = k.reshape(2, 4, 3, 8)
v = torch.linspace(3, 194, steps=192)
v = v.reshape(2, 4, 24)

k = k.permute(0, 3, 2, 1)
attn = torch.einsum(&quot;nchw,nwhu-&gt;nchu&quot;, q, k)

# Below is what doesn't work. I would like to get it such that hidden_states is a tensor of rank 2, 4, 24
hidden_states = torch.einsum(&quot;chw,whu-&gt;chu&quot;, attn, v)


Is there a permutation/transpose I could apply to q, k, v, or attn that would allow me to multiply into (2, 4, 24)?
I have yet to find one.
I currently receive this error: &quot;RuntimeError: einsum(): the number of subscripts in the equation (3) does not match the number of dimensions (4) for operand 0 and no ellipsis was given&quot; so I'm wondering how to use the ellipsis in this case, if that could be a solution.
Any explanation as to why this is or isn't possible would also be an excepted answer!

It seems like your q and k are 4D tensors of shape batch-channel-height-width (2x4x3x8). However, when considering attention mechanism, one disregard the spatial arrangement of the features and only treat them as a &quot;bag of features&quot;. That is, instead of q and k of shape 2x4x3x8 you should have 2x4x24:
q = torch.linspace(1, 192, steps=192)
q = q.reshape(2, 4, 3 * 8)  # collapse the spatial dimensions into a single one
k = torch.linspace(2, 193, steps=192)
k = k.reshape(2, 4, 3 * 8)  # collapse the spatial dimensions into a single one
v = torch.linspace(3, 194, steps=192)
v = v.reshape(2, 4, 24)

attn = torch.einsum(&quot;bcn,bcN-&gt;bnN&quot;, q, k)
# it is customary to convert the raw attn into probabilities using softmax
attn = torch.softmax(attn, dim=-1)
hidden_states = torch.einsum(&quot;bnN,bcN-&gt;bcn&quot;, attn, v)


