code: def run_fixed_mcmc_on_target(
      target: gym.targets.Model,
      init_x: jnp.ndarray,
      method: str,
      seed: jax.Array,
      num_warmup_steps: int,
      num_results: int,
      scalar_step_size: jnp.ndarray,
      vector_step_size: jnp.ndarray,
      num_integrator_steps: jnp.ndarray,
      damping: Optional[jnp.ndarray] = None,
  ):
      """Run a fixed-hyperparameter MCMC.
  
      Args:
        target: Target model.
        init_x: Initial constrained chain state.
        method: MCMC method. Either 'hmc' or 'malt'.
        seed: Random seed.
        num_warmup_steps: Number of warmup steps.
        num_results: Number of post-warmup results.
        scalar_step_size: Scalar step size.
        vector_step_size: Vector step size.
        num_integrator_steps: Number of integrator steps.
        damping: Damping.
  
      Returns:
         A tuple of final and traced results.
      """
      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(mcmc_state, seed):
          seed, jitter_seed, mcmc_seed = jax.random.split(seed, 3)
  
          if method == "hmc":
              cur_num_integrator_steps = jax.random.randint(
                  jitter_seed, [], 1, num_integrator_steps + 1
              )
              (
                  mcmc_state,
                  mcmc_extra,
              ) = fun_mc.hamiltonian_monte_carlo_step(
                  mcmc_state,
                  target_log_prob_fn=target_log_prob_fn,
                  step_size=scalar_step_size * vector_step_size,
                  num_integrator_steps=cur_num_integrator_steps,
                  seed=mcmc_seed,
              )
          elif method == "malt":
              cur_num_integrator_steps = num_integrator_steps
              (
                  mcmc_state,
                  mcmc_extra,
              ) = fun_mc.prefab.metropolis_adjusted_langevin_trajectories_step(
                  mcmc_state,
                  target_log_prob_fn=target_log_prob_fn,
                  step_size=scalar_step_size * vector_step_size,
                  momentum_refresh_fn=lambda m, seed: (  # pylint: disable=g-long-lambda
                      malt_lib._gaussian_momentum_refresh_fn(  # pylint: disable=protected-access
                          m,
                          damping=damping,
                          step_size=scalar_step_size / 2.0,
                          seed=seed,
                      )
                  ),
                  num_integrator_steps=cur_num_integrator_steps,
                  damping=damping,
                  seed=mcmc_seed,
              )
  
          traced = {
              "state": mcmc_state.state,
              "is_accepted": mcmc_extra.is_accepted,
              "log_accept_ratio": mcmc_extra.log_accept_ratio,
              "num_integrator_steps": cur_num_integrator_steps,
          }
  
          return (mcmc_state, seed), traced
  
      if method == "hmc":
          mcmc_state = fun_mc.hamiltonian_monte_carlo_init(
              state=init_z,
              target_log_prob_fn=target_log_prob_fn,
          )
      elif method == "malt":
          mcmc_state = fun_mc.prefab.metropolis_adjusted_langevin_trajectories_init(
              state=init_z,
              target_log_prob_fn=target_log_prob_fn,
          )
  
      _, trace = fun_mc.trace(
          (mcmc_state, seed), kernel, num_warmup_steps + num_results
      )
  
      warmed_up_state = trace["state"][num_warmup_steps:]
      cov = jnp.cov(
          warmed_up_state.reshape([-1, warmed_up_state.shape[-1]]),
          rowvar=False,
      )
      principal = jnp.linalg.eigh(cov)[1][:, -1]
  
      final = {
          "stats": compute_stats(
              state=warmed_up_state,
              num_grads=trace["num_integrator_steps"][
                  -num_warmup_steps:
              ].sum(),
              mean=warmed_up_state.mean((0, 1)),
              principal=principal,
          )
      }
      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 20 for calling a global function with example usage: run_fixed_mcmc_on_target(target, init_x, method, seed, num_warmup_steps, num_results, scalar_step_size, vector_step_size, num_integrator_steps, damping) and returns: trace,  final, target,  , mcmc_state,  seed,  traced