Numpy: multiply 3D array with 2D array

Question

I have 2 arrays, one of shape (4, 2, 2), another of shape (4, 2).

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)

So, for every 0<= i <= 3, I want to have np.dot(A[i, :, :], B[i, :]).

How to do it without using a for-loop?

Answer 1

(A @ B[...,np.newaxis])[...,0] should work (synactic sugar for np.matmul(A, B[...,np.newaxis])[...,0]).

Per the documentation for matmul:

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes.

You can also use Einstein summation np.einsum("ijk,ik->ij", A, B), although on my machine this is apparently slower than broadcasting matmul:

import timeit

import numpy as np

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)


def broadcast_matmul():
    return (A @ B[..., np.newaxis])[..., 0]


def list_map():
    return np.array(list(map(np.dot, A, B)))


def einsum():
    return np.einsum("ijk,ik->ij", A, B)


print(timeit.timeit(broadcast_matmul))
print(timeit.timeit(list_map))
print(timeit.timeit(einsum))

# Results on my machine:
# 1.870921753
# 6.705698656999999
# 2.765732645

Answer 2

Here's how this could be done using np.einsum:

np.einsum('ijk,ik->ij', A, B)

array([[0.39437083, 0.45360039],
       [0.75211742, 0.40669922],
       [0.18131254, 0.19473085],
       [0.2919673 , 0.10398859]])

---

Quick check using the first elements that are multiplied:

np.dot(A[0], B[0])
# array([0.39437083, 0.45360039])

Answer 3

This should do the trick: list(map(np.dot, A, B))