Batch matrix multiplication of 3D tensors

[KFrank]

Hi Shon!

@shon711

all_X  [btach_size, num_cats, k,     ffnn]
all_C  [btach_size, num_cats, ffnn,  ffnn]
all_Y  [btach_size, num_cats, k,     ffnn]

It’s not entirely clear to me what the values num_cats and k are.

for i in range(num_cat):
    x = all_X[:, i, :, :]                               # [btach_size, k,     ffnn]
    c = all_C[:, i, :, :]                               # [btach_size, ffnn,  ffnn]
    y = all_Y[:, i, :, :]                               # [btach_size, k,     ffnn]

    temp = torch.matmul(x, c)                           # [batch_size, k, ffnn_size]
    logits = torch.matmul(temp, y.permute([0, 2, 1])).  # [batch_size, k, k]

Let me assume that you want the shape of your final “all the multiplication
together” result to be [batch_size, num_cats, k, k].

You can compute your three-term product with einsum():

result = torch.einsum ('bnkf, bnfg, bnlg -> bnkl', all_X, all_C, all_Y)

Make sure you understand how einsum() works so that you are
contracting (“multiplying” together) the desired pairs of indices.

If all_C is constructed from the weights of Linears, you may need
to transpose the last two dimensions (the ffnn dimensions) of all_C,
either explicitly or by swapping the symbolic indices in the einsum()
expression, in order to replicate the result you would get by applying
the Linear to all_Y.

(This again ignores any bias terms that may have appeared in your
Linears.)

Best.

K. Frank