code: def run_meads_on_target(
      target: gym.targets.Model,
      num_adaptation_steps: int,
      num_results: int,
      thinning: int,
      seed: jax.Array,
      num_folds: int,
      num_chains: Optional[int] = None,
      init_x: Optional[jnp.ndarray] = None,
      ground_truth: Optional[Dict[str, Dict[str, jnp.ndarray]]] = None,
      save_warmup: bool = True,
      **meads_kwargs: Any,
  ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
      """Does a run of adaptive MCMC.
  
      Args:
        target: Target model.
        num_adaptation_steps: Number of adaptation steps.
        num_results: Number of post-adaptation results.
        thinning: How much to thin the chain by.
        seed: Random seed.
        num_folds: Number of folds.
        num_chains: Number of chains. Optional if `init_x` is specified.
        init_x: Initial constrained chain position. Optional if `num_chains` is
          specified.
        ground_truth: Ground truth means.
        save_warmup: Whether to save the warmup iterations.
        **meads_kwargs: Extra kwargs to pass to Adaptive MCMC.
  
      Returns:
         A tuple of final and traced results.
      """
      if init_x is None:
          init_x = get_init_x(target, num_chains)
      init_z = target.default_event_space_bijector.inverse(init_x)
  
      def target_log_prob_fn(x):
          return target.unnormalized_log_prob(x), ()
  
      target_log_prob_fn = fun_mc.transform_log_prob_fn(
          target_log_prob_fn, target.default_event_space_bijector
      )
  
      def kernel(meads_state, seed):
          seed, mcmc_seed = jax.random.split(seed)
  
          meads_state, meads_extra = meads_step(
              meads_state,
              target_log_prob_fn=target_log_prob_fn,
              seed=mcmc_seed,
              **meads_kwargs,
          )
  
          traced = {
              "state": meads_state.phmc_state.state.reshape(
                  [-1, init_z.shape[-1]]
              ),
              "is_accepted": meads_extra.phmc_extra.is_accepted.reshape(
                  [-1]
              ),
              "scalar_step_size": meads_extra.scalar_step_size.mean(0),
              "vector_step_size": meads_extra.vector_step_size.mean(
                  (0, 1)
              ),
              "max_eigenvalue": meads_extra.max_eigenvalue.mean(0),
              "damping": meads_extra.damping.mean(0),
              "extra": meads_extra.extra,
          }
  
          return (meads_state, seed), traced
  
      def chunked_kernel(*state):
          state, traced = fun_mc.trace(
              state, kernel, thinning, trace_mask=False
          )
          traced = traced.copy()
          traced["num_integrator_steps"] = thinning
          return state, traced
  
      seed, init_seed = jax.random.split(seed)
  
      meads_state = meads_init(
          init_z,
          target_log_prob_fn,
          num_folds=num_folds,
          seed=init_seed,
      )
  
      if save_warmup:
          _, trace = fun_mc.trace(
              (meads_state, seed),
              chunked_kernel,
              num_results + num_adaptation_steps,
          )
      else:
          state, _ = fun_mc.trace(
              (meads_state, seed),
              chunked_kernel,
              num_adaptation_steps,
              trace_fn=lambda *_: (),
          )
          _, trace = fun_mc.trace(state, chunked_kernel, num_results)
  
      warmed_up_steps = int(num_results * 0.8)
      warmed_up_state = trace["state"][-warmed_up_steps:]
  
      if ground_truth is None:
          cov = jnp.cov(
              warmed_up_state.reshape([-1, warmed_up_state.shape[-1]]),
              rowvar=False,
          )
          principal = jnp.linalg.eigh(cov)[1][:, -1]
          mean = warmed_up_state.mean((0, 1))
          bias = {}
      else:
          principal = ground_truth["principal"]
          mean = ground_truth["stats"]["m1"]["mean"]
          bias = {
              "bias": compute_bias(
                  trace["state"], ground_truth, mean, principal
              )
          }  # pytype: disable=wrong-arg-types  # jax-ndarray
  
      final = {
          "stats": compute_stats(
              state=warmed_up_state,
              num_grads=trace["num_integrator_steps"][
                  -warmed_up_steps:
              ].sum(),
              mean=mean,
              principal=principal,
          ),
          "final_x": target.default_event_space_bijector(
              trace["state"][-1]
          ),
          **bias,
      }
      return trace, final
  
imports: import functools
  import os
  from typing import Any, Callable, Dict, Optional, Mapping, NamedTuple, Tuple, Union
  from etils import epath
  import gin
  import immutabledict
  import jax
  import jax.numpy
  import numpy
  from discussion.adaptive_malt import utils
  from fun_mc import using_jax
  from inference_gym import using_jax
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 19 for calling a global function with example usage: run_meads_on_target(target, num_adaptation_steps, num_results, thinning, seed, num_folds, num_chains, init_x, ground_truth, save_warmup) and returns: trace,  final, target,  , meads_state,  seed,  traced, state,  traced