Batch matrix multiplication

1

Is there a way to do a batch matrix multiplication in pymc? The prototype operation in numpy
is:

B, M, N, P = 32, 100, 5, 2
w1_np = np.random.randn(B, M, N)
w2_np = np.random.randn(B, N, P)
data = w1_np.dot(w2_np) # shape: (B, M, P)

where I wanna carry out the B matrix multiplies (M, N), (N,P) → (M, P) in one shot and store them in
a 3D structure (B, M, P). My tentative pymc code is:

with pm.Model() as bmm:

    w1 = pm.Normal('w1', 0, sigma=1, shape=(B, M, N))
    w2 = pm.Normal('w2', 0, sigma=1, shape=(B, N, P))

    #ym = pm.math.dot(w1, w2)
    ym = w1 @ w2
    out = pm.Normal('out', mu=ym)#, observed=data)
    prior_data = pm.sample_prior_predictive()

but it does not work, while the version with ym = pm.math.dot(w1, w2) returns a result with different dimensions.

Thanks,
Marco

2

This is not what happens.

w1_np.dot(w2_np).shape == (B, M, B, P)

What you want is matmul:

np.matmul(w1_np, w2_np).shape == (B, M, P)

Which is what you get when you use the @ operator in numpy and in PyMC (but only the latest version, there was a mistake before).

To be safe you can use pytensor.tensor.matmul(w1, w2)

2 Likes
3

Thanks Ricardo,

Thanks for the reply. Updating to version 5.9, the @ operator in pymc works. Previously I had issues in 5.7.2

This is not what happens.

w1_np.dot(w2_np).shape == (B, M, B, P)

Right, it was a leftover of previous attempts I meant @ here.