Faster version of einsum 'ij,jlk->ilk'

def ij_jlk_to_ilk(A, B):
    """
    Faster version of einsum 'ij,jlk->ilk'
    """
    return A.dot(B.reshape(B.shape[0], -1)).reshape(A.shape[0], B.shape[1], B.shape[2])