Kronecker Product in PyTorch - with batch dimensions broadcast

kron.py

"""A part of the pylabyk library: numpytorch.py at https://github.com/yulkang/pylabyk"""
import torch
def kron(a, b):
    """
    Kronecker product of matrices a and b with leading batch dimensions.
    Batch dimensions are broadcast. The number of them mush
    :type a: torch.Tensor
    :type b: torch.Tensor
    :rtype: torch.Tensor
    """
    siz1 = torch.Size(torch.tensor(a.shape[-2:]) * torch.tensor(b.shape[-2:]))
    res = a.unsqueeze(-1).unsqueeze(-3) * b.unsqueeze(-2).unsqueeze(-4)
    siz0 = res.shape[:-4]
    return res.reshape(siz0 + siz1)

In response to : https://discuss.pytorch.org/t/kronecker-product/3919/11
Please let me know if you find any issues!

Note: this function is included in numpytorch.py, which is a part of pylabyk library I am maintaining.

This function as it is is not compatible with torch.compile. But a simple modification fixes it:

    siz1 = torch.Size([a.shape[-2] * b.shape[-2], a.shape[-1] * b.shape[-1]])