code: def compute_stats(
      state: jnp.ndarray,
      num_grads: jnp.ndarray,
      mean: jnp.ndarray,
      principal: jnp.ndarray,
  ) -> Dict[str, Dict[str, jnp.ndarray]]:
      """Computes MCMC statistics.
  
      Args:
        state: MCMC chain.
        num_grads: Number of gradient evaluations used to compute the chain.
        mean: State mean.
        principal: State principal component.
  
      Returns:
         A dict of statistics dicts.
      """
      transforms = _get_transforms(mean, principal)
  
      res = {}
      for name, transform_fn in transforms.items():
          transformed_state = transform_fn(state)
          ess = jax.lax.map(
              tfp.mcmc.effective_sample_size,
              jnp.transpose(transformed_state, [1, 0, 2]),
          ).mean(0)
          ess_per_step = ess / len(state)
          ess_per_grad = ess / num_grads
          rhat = tfp.mcmc.potential_scale_reduction(
              transformed_state, split_chains=True
          )
  
          white_state = transformed_state - transformed_state.mean(
              (0, 1)
          ) / transformed_state.std((0, 1))
          lag_1_autocorr = (white_state[:-1] * white_state[1:]).sum(
              (0, 1)
          )
  
          stats = {}
          stats["mean"] = transformed_state.mean((0, 1))
          stats["ess"] = ess
          stats["ess_per_grad"] = ess_per_grad
          stats["ess_per_step"] = ess_per_step
          stats["rhat"] = rhat
          stats["lag_1_autocorr"] = lag_1_autocorr
          res[name] = stats
  
      return res
  
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 11 for calling a global function with example usage: compute_stats(state, num_grads, mean, principal) and returns: res