Vectorization of Product Over

November 5, 2021

I am seeking a vectorized form of the following computation:

import numpy as np
D = 100
N = 1000
K = 10

X = np.random.uniform(0, 1, (K, N))
T = np.random.uniform(0, 1000, (D, N))
out = np.zeros((D, K))

for i in range(D):
    for j in range(K):
        out[i, j] = np.prod(X[j, :] ** T[i, :])

There are einsum-style things I’ve tried, but the presence of the np.prod is throwing me off a bit.

EDIT: Reduced size of matrices.

>Solution :

I’m trying to make the broadcasting as explicit as possible – the None introduces an additional dummy dimension of size 1:

out = np.prod(X[None, :, :] ** T[:, None, :], axis=2)

It is easy to see how it works if we recall the shapes: X.shape = (K, N), T.shape = (D, N) and out.shape = (D, K). With the dummy dimension we basically take something of (1, K, N) to the power of (D, 1, N) which results in (D, K, N). Finally if we reduce via product over the last dimension we get our desired output of (D, K).