Project: s41467-024-53672-w | Function: infidelity_chunked (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes the infidelity between variational and target quantum states using chunked Monte Carlo sampling, which is central to the variational optimization process

**KEY TERMS:** {
  "vstate": "variational state |ψ⟩",
  "op": "infidelity operator I_op",
  "chunk_size": "Monte Carlo sampling chunk size",
  "L_stats": "infidelity statistics",
  "sigma": "configuration samples σ"
}

**CONTEXT CODE:**     log_phi_y = apply_fun_phi(pars_phi, y)

    # compute the weights
    if include_sqrt:
        log_w += log_phi_y.real + log_psi_x.real

    if include_R:
        log_Uphi_y = _log_Upsi(apply_fun_phi, pars_phi, y, O, *O_args)
        log_Rloc = log_Uphi_y - log_phi_y
        log_w += 2 * log_Rloc.real

    return log_w


def _local_kernel_log_fidelity_jax(  # NO weights
    O: ContinuousOperator,
    O_dagger: ContinuousOperator,
    apply_fun_psi: Callable,
    apply_fun_phi: Callable,
    pars_psi: PyTree,
    pars_phi: PyTree,
    x: Array,
    y: Array,
    O_args: Any,
    O_dagger_args: Any,
    cv_coeff: float = None,
):
    # x samples
    log_Uphi_x = _log_Upsi(apply_fun_phi, pars_phi, x, O, *O_args)
    log_psi_x = apply_fun_psi(pars_psi, x)
    # y samples
    log_psi_y = apply_fun_psi(pars_psi, y)
    log_Uphi_y = _log_Upsi(apply_fun_phi, pars_phi, y, O, *O_args)

    log_Floc_xy = (log_Uphi_x - log_psi_x) + (log_psi_y - log_Uphi_y)

    log_res = log_Floc_xy
    if cv_coeff is not None:
        log_res = logsumexp(
            jnp.stack(
                [log_res, 2 * log_Floc_xy.real, jnp.zeros_like(log_res)], axis=-1
            ),
            b=jnp.array([1, cv_coeff, -cv_coeff]),
            axis=-1,
        )

    return log_res


@expect.dispatch
def infidelity_chunked(vstate: MCState, op: InfidelityOperatorUPsi, chunk_size: int):
    raise NotImplementedError("Function implementation has been masked")


@expect_and_forces.dispatch
def infidelity_and_forces_chunked(  # noqa: F811
    vstate: MCState,
    op: InfidelityOperatorUPsi,
    chunk_size: int,
    *,
    mutable,
):
    L_stats, L_grad, wc = _infidelity_and_grad_chunked(
        vstate, op, chunk_size, return_grad=True
    )
    op._wc = wc
    return L_stats, L_grad


def _infidelity_and_grad_chunked(
    vstate: MCState,
    op: InfidelityOperatorUPsi,
    chunk_size: int,
    return_grad: bool = True,
):
    # now in a single function
    if op.hilbert != vstate.hilbert:
        raise TypeError("Hilbert spaces should match")

    validate_arguments(vstate, op, chunk_size=chunk_size)

    (
        sigma,
        sigma_t,
        args,
        args_in_axes,
    ) = get_local_kernel_arguments_infidelity_with_in_axes(
        vstate, op.target, op._U, op._U_dagger
    )

    # functions to compute Floc etc
    local_log_infidelity_kernel = nk.utils.HashablePartial(
        _local_kernel_log_fidelity_jax, op._U, op._U_dagger
    )
    local_log_weights_kernel = nk.utils.HashablePartial(
        _local_kernel_log_weights_jax, op._U, op._U_dagger, include_sqrt=op._sample_sqrt
    )

    return infidelity_sampling_MCState(
        chunk_size,
        vstate._apply_fun,
        op.target._apply_fun,
        vstate.parameters,
        op.target.parameters,
        vstate.model_state,
        op.target.model_state,
        local_log_infidelity_kernel,
        local_log_weights_kernel,
        sigma,
        sigma_t,
        args,
        args_in_axes,
        op.cv_coeff,
        return_grad=return_grad,
    )


def local_log_weights_kernel_wrapper(  # add some spice
    local_log_weights_kernel,
    afun,
    params,
    σ,
    afun_t,
    params_t,
    σ_t,
    args,
    model_state=None,
    model_state_t=None,
):
    """Part handling the local fidelity kernel."""

    # to be sure, we block some gradients (speeds some things up)

    model_state = model_state or {}
    model_state_t = model_state_t or {}

    W = {"params": params, **model_state}
    W_t = {"params": params_t, **model_state_t}

    batch_shape = σ.shape[:-1]
    σ = σ.reshape(-1, σ.shape[-1])
    σ_t = σ_t.reshape(-1, σ_t.shape[-1])

    log_w = local_log_weights_kernel(afun, afun_t, W, W_t, σ, σ_t, *args)
    log_w = log_w.reshape(*batch_shape)
    return log_w


def local_log_fidelity_kernel_wrapper(  # add some spice
    local_log_fidelity_kernel,


Complete the function implementation replacing "raise NotImplementedError".