Kronnecker product on just two dimensions of 3 dimensional arrays

Question

Suppose two matrices:

a = np.array([[1,2],
              [3,4]])
b = np.array([[5,6],
              [7,8]])

which would give Kronecker product ab = np.kron(a,b):

array([[ 1,  2,  2,  4],
       [ 3,  4,  6,  8],
       [ 3,  6,  4,  8],
       [ 9, 12, 12, 16]])

Now suppose there are three copies of these matrices in two arrays, like this:

c = np.stack([a,a,a])
d = np.stack([b,b,b])

I want to compute the Kronecker product of c and d such that the output is a 3 index array corresponding to 3 copies of ab, i.e. with shape (3,4,4). However, simply performing kron(c,d) outputs shape (9,4,4), which has more entries than needed and cannot be reshaped appropriately. Could you please help understand how to do this?

Answer 1

This works, assuming you have no problem with the intermediate list (numpy needs the size before allocating the array):

import numpy as np

a = np.array([[1, 2],
              [3, 4]])
b = np.array([[5, 6],
              [7, 8]])

c = np.stack([a, a, a])
d = np.stack([b, b, b])

result = np.array(list(np.kron(x, y) for x, y in zip(c, d)))

print(result)
print(result.shape)

Output:

[[[ 5  6 10 12]
  [ 7  8 14 16]
  [15 18 20 24]
  [21 24 28 32]]

 [[ 5  6 10 12]
  [ 7  8 14 16]
  [15 18 20 24]
  [21 24 28 32]]

 [[ 5  6 10 12]
  [ 7  8 14 16]
  [15 18 20 24]
  [21 24 28 32]]]
(3, 4, 4)

Answer 2

res=np.zeros((3,4,4))
res[:] = np.kron(a,b)

Should work, broadcasting the kron to all 3 planes.

kron is specialized rearrangement of the outer product of a and b. A (2,2,2,2) rearranged to (4,4). I worked out the details in another post:

Why is numpy's kron so fast?

Your (3,4,4) could be obtained from a (3,2,2,2,2), but it's not standard, so there isn't an out-of-the-box function. You could try adapting my answer.

---

In [246]: a = np.array([[1,2],
     ...:               [3,4]])
     ...: b = np.array([[5,6],
     ...:               [7,8]])

In [249]: np.kron(a,b)
Out[249]: 
array([[ 5,  6, 10, 12],
       [ 7,  8, 14, 16],
       [15, 18, 20, 24],
       [21, 24, 28, 32]])

As I showed before, kron can be produced by applying a transpose and shape to an outer product. We can use einsum for both the outer and transpose:

In [253]: np.einsum('ij,kl->ikjl',a,b)     # ikjl instead ijkl
Out[253]: 
array([[[[ 5,  6],
         [10, 12]],

        [[ 7,  8],
         [14, 16]]],


       [[[15, 18],
         [20, 24]],

        [[21, 24],
         [28, 32]]]])
In [254]: np.einsum('ij,kl->ikjl',a,b).reshape(4,4)
Out[254]: 
array([[ 5,  6, 10, 12],
       [ 7,  8, 14, 16],
       [15, 18, 20, 24],
       [21, 24, 28, 32]])

Generalizing that to arrays (3,2,2) shape, we can add an extra 'batch' dimension:

In [255]: c = np.stack([a,a,a])
     ...: d = np.stack([b,b,b])
In [256]: c
Out[256]: 
array([[[1, 2],
        [3, 4]],

       [[1, 2],
        [3, 4]],

       [[1, 2],
        [3, 4]]])
In [257]: c.shape
Out[257]: (3, 2, 2)
In [258]: np.einsum('aij,akl->aikjl',c,d).reshape(3,4,4)
Out[258]: 
array([[[ 5,  6, 10, 12],
        [ 7,  8, 14, 16],
        [15, 18, 20, 24],
        [21, 24, 28, 32]],

       [[ 5,  6, 10, 12],
        [ 7,  8, 14, 16],
        [15, 18, 20, 24],
        [21, 24, 28, 32]],

       [[ 5,  6, 10, 12],
        [ 7,  8, 14, 16],
        [15, 18, 20, 24],
        [21, 24, 28, 32]]])

But if we know that c and d are just replications of a and b, the broadcasted solution is fastester

In [260]: res = np.zeros((3,4,4),int)
In [261]: res[:] = np.kron(a,b)

or even better (a view of the kron witout copies):

np.broadcast_to(np.kron(a,b),(3,4,4))

Some timings:

In [280]: timeit np.einsum('aij,akl->aikjl',c,d).reshape(3,4,4)
10.2 µs ± 22.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [281]: timeit res=np.zeros((3,4,4),int);res[:] = np.kron(a,b)
47.5 µs ± 1.66 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [282]: timeit np.broadcast_to(np.kron(a,b),(3,4,4))
57.6 µs ± 1.76 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [283]: timeit np.array(list(np.kron(x, y) for x, y in zip(c, d)))
143 µs ± 319 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

I'm bit surprised how much faster the einsum is. Also that braodcast_to isn't faster, though that's largely the fault of kron (which my previous answer showed was on the slow side).

Answer 3

My solution for this would go like this

import numpy as np

def kron(a, b):

    ai = a[0]
    bi = b[0]

    b = np.asanyarray(b)
    a = np.array(a, copy=False, subok=True, ndmin=b.ndim)

    ndb, nda = bi.ndim, ai.ndim

    b_size = a.shape[0]

    as_ = ai.shape
    bs = bi.shape

    nd = ndb
    if (ndb != nda):
        if (ndb > nda):
            as_ = (1,) * (ndb - nda) + as_
        else:
            bs = (1,) * (nda - ndb) + bs
            nd = nda

    res = np.einsum("ij,ik->ijk", a.reshape((a.shape[0], -1)), b.reshape((b.shape[0], -1)))
    res = res.reshape((b_size, )+(as_ + bs))
    axis = nd - 1
    result = []
    for i in range(b_size):
        r = res[i]
        for _ in range(nd):
            r = np.concatenate(r, axis=axis)
        result.append(r)
    return np.array(result)

x = np.array([[1,2],
              [3,4]])
y = np.array([[5,6],
              [7,8]])

c = np.stack([x,x,x])
d = np.stack([y,y,y])

k = kron(c, d)
result = np.array(list(np.kron(x, y) for x, y in zip(c, d)))

print(np.allclose(k, result))

As @hpaulj mentioned in their answer, there is no built-in function to do this, and you would have to do some non-trivial reshape operation on the outer product to get this working.

The solution has been modified based on the implementation of kron in numpy.

I basically just added another outer loop on top of the loop that the numpy implementation is using to do the concatenation operation. This is meant to handle each outer product between the elements of c and d separately. Everything else is pretty much the same.

Note that this is not a generalisation of the kron implementation of numpy because this is not going to work for cases where you just want to do a straightforward kron operation. At least not without some additional modifications. I am also ignoring some edge cases that numpy is handling for clarity. However, this should work for your use-case.