Combining Tensors in a Simple Multiplicative Way

Hello, my first post here. I did search for the question but wasn’t able to find it elsewhere in the forum.

I have a tensor X of dimension A, and a tensor Y of dimensions MxNxP and want to create a new tensor Z of dimensions MxAxNxP where Z[m,a,n,p] = X[a]*Y[m,n,p].

I did find a solution but was rather inelegant and possibly inefficient.

There is a simple way of doing this?

Broadcasting should work as seen here:

# setup
A, M, N, P = 2, 3, 4, 5
X = torch.randn(A)
Y = torch.randn(M, N, P)
Z_ref = torch.zeros(M, A, N, P)

# slow approach to create reference
for m in range(M):
    for a in range(A):
        for n in range(N):
            for p in range(P):
                Z_ref[m, a, n, p] = X[a] * Y[m, n, p]

# non-loop approach
Z = X[None, :, None, None] * Y[:, None, ...]

# compare
print((Z_ref - Z).abs().max())
# tensor(0.)
1 Like

Thanks @ptrblck, that is a very nice solution!