How can I use vectorize or array broadcasting to get rid of this for-loop?

Question

Is there a way to use np.vectorize or numpy array broadcasting to get rid of this for-loop?

xs = np.random.uniform(-2,2,size=(3,2))

Qs = np.array([
    [
        [4,2],
        [2,2]
    ],
    [
        [1,1],
        [1,1]
    ],
    [
        [2,2],
        [2,2]
    ]
])
rs = np.array([
    [1,-1],
    [1,1],
    [-1,1]
])

g = lambda Q, r, x: 0.5*x.T @ Q @ x + r @ x

# loop over all of the elements of Qs, rs, and xs 
# evaluating g at each point
results = []
for Q, r, x in zip(Qs, rs, xs):
    results.append(g(Q, r, x))
print(np.array(results))

I've tried just using array-broadcasting:

results = g(Qs, rs, xs)

but it states

ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 3)

and I have tried vectorizing the funciton:

g_vec = np.vectorize(g)
results = g_vec(Qs, rs, xs)
print(results)

but it gives this error:

TypeError: unsupported operand type(s) for @: 'numpy.float64' and 'numpy.int32'

and doesn't seem to loop over each Q, r and x np.array

Answer 1

You need to add a signature to your vectorization function to properly multiply the matrices:

g_vec = np.vectorize(g, signature='(n,n),(n),(n)->()')
results = g_vec(Qs, rs, xs)
print(results)

More on signatures.