code: def test_idata_contains_stats(sampler_name: str):
      """Tests whether sampler statistics were written to sample_stats
      group of InferenceData"""
      if sampler_name == "sample_blackjax_nuts":
          sampler = sample_blackjax_nuts
      elif sampler_name == "sample_numpyro_nuts":
          sampler = sample_numpyro_nuts
  
      with pm.Model():
          pm.Normal("a")
          idata = sampler(tune=50, draws=50)
  
      stats = idata.get("sample_stats")
      assert stats is not None
      n_chains = stats.sizes["chain"]
      n_draws = stats.sizes["draw"]
  
      # Stats vars expected for both samplers
      expected_stat_vars = {
          "acceptance_rate": (n_chains, n_draws),
          "diverging": (n_chains, n_draws),
          "energy": (n_chains, n_draws),
          "tree_depth": (n_chains, n_draws),
          "lp": (n_chains, n_draws),
      }
      # Stats only expected for blackjax nuts
      if sampler_name == "sample_blackjax_nuts":
          blackjax_special_vars = {}
          stat_vars = expected_stat_vars | blackjax_special_vars
      # Stats only expected for numpyro nuts
      elif sampler_name == "sample_numpyro_nuts":
          numpyro_special_vars = {
              "step_size": (n_chains, n_draws),
              "n_steps": (n_chains, n_draws),
          }
          stat_vars = expected_stat_vars | numpyro_special_vars
      # test existence and dimensionality
      for stat_var, stat_var_dims in stat_vars.items():
          assert stat_var in stats.variables
          assert stats.get(stat_var).values.shape == stat_var_dims
  
imports: import warnings
  from typing import Any, Callable, Dict, Optional
  from unittest import mock
  import arviz
  import jax
  import jax.numpy
  import numpy
  import pytensor
  import pytensor.tensor
  import pytest
  from numpyro.infer import MCMC
  from pytensor.compile import SharedVariable
  from pytensor.graph import graph_inputs
  import pymc
  from pymc import ImputationWarning
  from pymc.distributions.multivariate import PosDefMatrix
  from pymc.sampling.jax import _get_batched_jittered_initial_points, _get_log_likelihood, _numpyro_nuts_defaults, _replace_shared_variables, _update_numpyro_nuts_kwargs, get_jaxified_graph, get_jaxified_logp, sample_blackjax_nuts, sample_numpyro_nuts
label_desc: Use image node_id 15 for calling a global function with example usage: test_idata_contains_stats(sampler_name) without return types