Vectorization of Product Over

Question

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.

Can you make a more minimal example? I don't have 80gb of memory lying around. — user3483203, Nov 04 '21 at 13:31
Your code just results in all zeros for me. Are these actual inputs? — user3483203, Nov 04 '21 at 13:36
That's fine; as long as the expression is correct. The minimal example is for functionality, not utility (I can't upload my original matrices here) — rw435, Nov 04 '21 at 13:39
@user3483203. You can decrease any of the sizes without harming the concept... — Mad Physicist, Nov 04 '21 at 13:39

score 2 · Accepted Answer · answered Nov 04 '21 at 13:51

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).

Vectorization of Product Over

1 Answers1