sprk
sprk

Reputation: 13

Pytorch how to multiply tensors of variable size except the first dimention

I have one tensor which is A = 40x1.

i need to multiply this one with 3 other tensors: B = 40x100x384, C = 40x10, D=40x10.

for example in tensor B, we got 40 100x384 matrixes and i need each one of these matrixes to be multiplied with its corresponding element from A

what is the best way to do this in pytorch? Suppose that we could have more matrixes like B,C,D they will always be in the style 40xKxL or 40xJ

Upvotes: 1

Views: 511

Answers (1)

iacolippo
iacolippo

Reputation: 4513

If I understand correctly, you want to multiply every i-th matrix K x L by the corresponding i-th scalar in A.

One possible way is:

(A * B.view(len(A), -1)).view(B.shape)

Or you can use the power of broadcasting:

A = A.reshape(len(A), 1, 1)
# now A is (40, 1, 1) and you can do
A*B
A*C
A*D

essentially each trailing dimension equal to 1 in A is stretched and copied to match the other matrix.

Upvotes: 1

Related Questions