def _kp(a, b):
    """Special case Kronecker tensor product of a[i] and b[i] at each
    time interval i for i = 0 .. N-1
    It is specialized for the case where both a and b are shape N x m x 1
    """
    if a.shape != b.shape or a.shape[-1] != 1:
        raise(ValueError)
    N = a.shape[0]
    # take the outer product over the last two axes, then reshape:
    return np.einsum('ijk,ilk->ijkl', a, b).reshape(N, -1, 1)