Batch Kronecker product of tensors

Question

I have two tensors that are batches of matrices:

x = torch.randn(100,10,10) 
y = torch.randn(100,2,2)

I want to parallelize the kronecker on each matrix, not doing the kronecker product of the tensors. torch.kron(x,y) gives me a tensor of size (10000, 60,60), but I want a shape out of size (100, 60, 60) that computes the kronecker product for each matrix. Is there any way to do so ?

Something like torch.kron(x,y, start_dim = 1) is what I tried but it does not seem to be implemented. (I want to do this in torch for R but something that works in python would already be okay)

Thanks

Answer 1

x = torch.randn(100,10,10) 
y = torch.randn(100,2,2)
z = torch.einsum('iab,icd->iacbd', x, y).view((100, 20, 20))

which is verified below as equivalent to taking Kronecker product over the last two dimensions:

for xx, yy, zz in zip(x, y, z):
    assert zz.allclose(torch.kron(xx, yy))