numpy kron along a given axis… here is a solution to the problem.
numpy kron along a given axis
Is there a function that applies the Kronecker product along a given axis? For example, given a two-dimensional array of shapes a.shape == (n, k) and b.shape = (n, l), calculate c of shape
c.shape == (n, k
*l),
and the result is equivalent to:
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])
Solution
There is no built-in, but we can use outer elementwise-multiplication
to keep 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)