import torch


def kron_einsum_batched_1D(A: torch.Tensor, B: torch.Tensor): 
    """ 
    Batched Version of Kronecker Products 
    :param A: has shape (b, c, a) 
    :param B: has shape (b, c, k) 
    :return: (b, c, ak) 
    """   
    res = torch.einsum('ba,bk->bak', A, B).view(A.size(0), A.size(1)*B.size(1) ) 
#     res = torch.einsum('bca,bck->bcak', A, B).view(A.size(0), A.size(1), A.size(2)*B.size(2) ) 
    return res