code: def _sample_stats_to_xarray(posterior):
      """Extract sample_stats from NumPyro posterior."""
      rename_key = {
          "potential_energy": "lp",
          "adapt_state.step_size": "step_size",
          "num_steps": "n_steps",
          "accept_prob": "acceptance_rate",
      }
      data = {}
      for stat, value in posterior.get_extra_fields(
          group_by_chain=True
      ).items():
          if isinstance(value, (dict, tuple)):
              continue
          name = rename_key.get(stat, stat)
          value = value.copy()
          data[name] = value
          if stat == "num_steps":
              data["tree_depth"] = np.log2(value).astype(int) + 1
      return data
  
imports: import logging
  import os
  import re
  from datetime import datetime
  from functools import partial
  from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Union
  import arviz
  import jax
  import jax.numpy
  import numpy
  import pytensor.tensor
  from arviz.data.base import make_attrs
  from jax.lax import scan
  from pytensor.compile import SharedVariable, Supervisor, mode
  from pytensor.graph.basic import graph_inputs
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.replace import clone_replace
  from pytensor.link.jax.dispatch import jax_funcify
  from pytensor.raise_op import Assert
  from pytensor.tensor import TensorVariable
  from pytensor.tensor.random.type import RandomType
  from pymc import Model, modelcontext
  from pymc.backends.arviz import coords_and_dims_for_inferencedata, find_constants, find_observations
  from pymc.distributions.multivariate import PosDefMatrix
  from pymc.initial_point import StartDict
  from pymc.logprob.utils import CheckParameterValue
  from pymc.sampling.mcmc import _init_jitter
  from pymc.util import RandomSeed, RandomState, _get_seeds_per_chain, get_default_varnames
label_desc: Use image node_id 6 for calling a global function with example usage: _sample_stats_to_xarray(posterior) and returns: data