import torch


def kron(A, B):
    """
    Kronecker Product.
    Works with batch dimemsion(s) - requires both A and B have same batch dims.
    """
    A_shp, B_shp = A.shape, B.shape
    assert A_shp[:-2] == B_shp[:-2]
    kron_block = torch.matmul(
        A[..., :, None, :, None], B[..., None, :, None, :]
    )
    kron = kron_block.reshape(
        A_shp[:-2] + (A_shp[-2] * B_shp[-2], A_shp[-1] * B_shp[-1],)
    )
    return kron