numpy kron along a given axis

Question

Is there a function for applying the Kronecker product along a given axis? For example, given 2D arrays a and b of shapes a.shape == (n, k) and b.shape == (n, l), compute c of shape c.shape == (n, k*l) with the equivalent result of:

c = np.empty((a.shape[0], a.shape[1] * b.shape[1]))
for i in range(c.shape[0]):
    c[i,:] = np.kron(a[i], b[i])

Answer 1

There isn't a built-in, but we can use outer elementwise-multiplication keeping their first axis aligned and then reshape -

c = (a[:,:,None]*b[:,None,:]).reshape(a.shape[0],-1)

Alternatively, we can use einsum -

c = np.einsum('nk,nl->nkl',a,b).reshape(a.shape[0],-1)