def METHOD_NAME(data, batch_size):
    data = jnp.atleast_1d(data)
    if data.ndim == 1:
        data = data.reshape((1, -1))
    if data.ndim > 2:
        raise NotImplementedError("Statistics are implemented only for ndim<=2")
    mean = _mean(data)
    variance = _var(data)
    ts = _total_size(data)
    bare_var = variance
    batch_var, n_batches = _batch_variance(data)
    l_block = max(1, data.shape[1] // batch_size)
    block_var, n_blocks = _block_variance(data, l_block)
    tau_batch = ((ts / n_batches) * batch_var / bare_var - 1) * 0.5
    tau_block = ((ts / n_blocks) * block_var / bare_var - 1) * 0.5
    batch_good = (tau_batch < 6 * data.shape[1]) * (n_batches >= batch_size)
    block_good = (tau_block < 6 * l_block) * (n_blocks >= batch_size)
    stat_dtype = nkjax.dtype_real(data.dtype)
    # if batch_good:
    #    error_of_mean = jnp.sqrt(batch_var / n_batches)
    #    tau_corr = jnp.max(0, tau_batch)
    # elif block_good:
    #    error_of_mean = jnp.sqrt(block_var / n_blocks)
    #    tau_corr = jnp.max(0, tau_block)
    # else:
    #    error_of_mean = jnp.nan
    #    tau_corr = jnp.nan
    # jax style
    def batch_good_err(args):
        batch_var, tau_batch, *_ = args
        error_of_mean = jnp.sqrt(batch_var / n_batches)
        tau_corr = jnp.clip(tau_batch, 0)
        return jnp.asarray(error_of_mean, dtype=stat_dtype), jnp.asarray(
            tau_corr, dtype=stat_dtype
        )
    def block_good_err(args):
        _, _, block_var, tau_block = args
        error_of_mean = jnp.sqrt(block_var / n_blocks)
        tau_corr = jnp.clip(tau_block, 0)
        return jnp.asarray(error_of_mean, dtype=stat_dtype), jnp.asarray(
            tau_corr, dtype=stat_dtype
        )
    def nan_err(args):
        return jnp.asarray(jnp.nan, dtype=stat_dtype), jnp.asarray(
            jnp.nan, dtype=stat_dtype
        )
    def batch_not_good(args):
        batch_var, tau_batch, block_var, tau_block, block_good = args
        return jax.lax.cond(
            block_good,
            block_good_err,
            nan_err,
            (batch_var, tau_batch, block_var, tau_block),
        )
    error_of_mean, tau_corr = jax.lax.cond(
        batch_good,
        batch_good_err,
        batch_not_good,
        (batch_var, tau_batch, block_var, tau_block, block_good),
    )
    if n_batches > 1:
        N = data.shape[-1]
        if not config.netket_use_plain_rhat:
            # compute split-chain batch variance
            local_batch_size = data.shape[0]
            if N % 2 == 0:
                # split each chain in the middle,
                # like [[1 2 3 4]] -> [[1 2][3 4]]
                batch_var, _ = _batch_variance(
                    data.reshape(2 * local_batch_size, N // 2)
                )
            else:
                # drop the last sample of each chain for an even split,
                # like [[1 2 3 4 5]] -> [[1 2][3 4]]
                batch_var, _ = _batch_variance(
                    data[:, :-1].reshape(2 * local_batch_size, N // 2)
                )
        # V_loc = _np.var(data, axis=-1, ddof=0)
        # W_loc = _np.mean(V_loc)
        # W = _mean(W_loc)
        # # This approximation seems to hold well enough for larger n_samples
        W = variance
        R_hat = jnp.sqrt((N - 1) / N + batch_var / W)
    else:
        R_hat = jnp.nan
    res = Stats(mean, error_of_mean, variance, tau_corr, R_hat)
    return res