How to use the thewalrus._hermite_multidimensional.hafnian_batched function in thewalrus

pref = np.exp(prefexp.conj())
    if post_select is None:
        choi_r = kwargs.get("choi_r", None)
        if choi_r is None:
            denom = np.sqrt(np.sqrt(np.linalg.det(Q).real))
        else:
            rescaling = np.concatenate(
                [np.ones([N // 2]), (1.0 / np.tanh(choi_r)) * np.ones([N // 2])]
            )
            B = np.diag(rescaling) @ B @ np.diag(rescaling)
            gamma = rescaling * gamma
            denom = np.sqrt(np.sqrt(np.linalg.det(Q / np.cosh(choi_r)).real))

        psi = (
            np.real_if_close(pref)
            * hafnian_batched(B.conj(), cutoff, mu=gamma.conj(), renorm=True)
            / denom
        )
    else:
        M = N - len(post_select)
        psi = np.zeros([cutoff] * (M), dtype=np.complex128)

        for idx in product(range(cutoff), repeat=M):
            el = []

            counter = count(0)
            modes = (np.arange(N)).tolist()
            el = [post_select[i] if i in post_select else idx[next(counter)] for i in modes]
            psi[idx] = pure_state_amplitude(
                mu, cov, el, check_purity=False, include_prefactor=False, hbar=hbar
            )