Batch 3DTensor-Matrix Multiplication?

Hi there,

I use bmm to multiply batches of matrices:

data = Variable(FloatTensor(BATCHSIZE,42,100).normal_(std=1)) # BATCH x 42 x 100
A = Variable(FloatTensor(BATCHSIZE,100,100).normal_(std=1))
B = bmm(data,A) # BATCH x 42 x 100 (OK)

Now I’d like to batch-multiply data with multiple matrices:

A = Variable(FloatTensor(BATCHSIZE,NLAYERS,100,100).normal_(std=1))
B = bmm(data,A) # doesn’t work

bmm() gives me the following error:

RuntimeError: invalid argument 2: expected 3D tensor, got 4D

Is there a way to batch-multiple matrices (data) with 3D-Tensors (A) ?

Thanks in advance!

1 Like

Hi,

You can use .view() on your 4D input to make it a 3D tensor with a batch size of BATCHSIZE * NLAYERS:

A = Variable(FloatTensor(BATCHSIZE,NLAYERS,100,100).normal_(std=1))
A = A.view(BATCHSIZE * NLAYERS, 100, 100)
B = bmm(data,A) # BATCHSIZE * NLAYERS x 42 x 100
B = B.view(BATCHSIZE, NLAYERS, 42, 100) # BATCHSIZE x NLAYERS x 42 x 100
2 Likes

That’s exactly what I was looking for - cool, thx!