Matrix Vector Product across Multiple Dimensions

Question

I have two arrays:

A = torch.rand((64, 128, 10, 10))
B = torch.rand((64, 128, 10))

I would like to compute the product, represented by C, where we do a matrix-vector multiplication across the first and second dimensions of A and B, so:

# C should have shape: (64, 128, 10)
for i in range(0, 64):
   for j in range(0, 128):
       C[i,j] = torch.matmul(A[i,j], B[i,j])

Does anyone know how to do this using torch.einsum? I tried the following, but I am getting an incorrect result.

C = torch.einsum('ijkl, ijk -> ijk', A, B)

Answer 1

Here's the options with numpy. (I don't have torch)

In [120]: A = np.random.random((64, 128, 10, 10))
     ...: B = np.random.random((64, 128, 10))

Your iterative reference case:

In [122]: C = np.zeros((64,128,10))
     ...: # C should have shape: (64, 128, 10)
     ...: for i in range(0, 64):
     ...:    for j in range(0, 128):
     ...:        C[i,j] = np.matmul(A[i,j], B[i,j])
     ...: 

matmul with full broadcasting:

In [123]: D  = np.matmul(A, B[:,:,:,None])
In [125]: C.shape
Out[125]: (64, 128, 10)
In [126]: D.shape            # D has an extra size 1 dimension
Out[126]: (64, 128, 10, 1)
In [127]: np.allclose(C,D[...,0])    # or use squeeze
Out[127]: True

The einsum equivalent:

In [128]: E = np.einsum('ijkl,ijl->ijk', A, B)
In [129]: np.allclose(C,E)
Out[129]: True