code: def adaptive_hamiltonian_monte_carlo_init(
      state: fun_mc.ArrayNest,
      target_log_prob_fn: fun_mc.PotentialFn,
      step_size: float | fun_mc.FloatArray = 1e-2,
      initial_mean: float | fun_mc.FloatNest = 0.0,
      initial_scale: float | fun_mc.FloatNest = 1.0,
      scale_smoothing_steps: int | fun_mc.IntArray = 10,
  ) -> AdaptiveHamiltonianMonteCarloState:
      """Initializes `AdaptiveHamiltonianMonteCarloState`.
  
      Args:
        state: Initial state of the chain.
        target_log_prob_fn: Target log prob fn.
        step_size: Initial scalar step size.
        initial_mean: Initial mean for computing the running variance estimate. Must
          broadcast structurally and tensor-wise with state.
        initial_scale: Initial scale for computing the running variance estimate.
          Must broadcast structurally and tensor-wise with state.
        scale_smoothing_steps: How much weight to assign to the `initial_mean` and
          `initial_scale`. Increase this to stabilize early adaptation.
  
      Returns:
        adaptive_hmc_state: State of the `adaptive_hamiltonian_monte_carlo_step`
          `TransitionOperator`.
      """
      hmc_state = fun_mc.hamiltonian_monte_carlo_init(
          state, target_log_prob_fn
      )
      dtype = util.flatten_tree(hmc_state.state)[0].dtype
      chain_ndims = len(hmc_state.target_log_prob.shape)
      running_var_state = fun_mc.running_variance_init(  # pytype: disable=wrong-keyword-args
          shape=util.map_tree(
              lambda s: s.shape[chain_ndims:], hmc_state.state
          ),
          dtype=util.map_tree(lambda s: s.dtype, hmc_state.state),
      )
      initial_mean = fun_mc.maybe_broadcast_structure(
          initial_mean, state
      )
      initial_scale = fun_mc.maybe_broadcast_structure(
          initial_scale, state
      )
  
      # It's important to add some smoothing here, as initial updates can be very
      # different than the stationary distribution.
      # TODO(siege): Add a pseudo-update functionality to avoid fiddling with the
      # internals here.
      running_var_state = running_var_state._replace(
          num_points=util.map_tree(
              lambda p: (  # pylint: disable=g-long-lambda
                  int(np.prod(hmc_state.target_log_prob.shape))
                  * jnp.array(scale_smoothing_steps, p.dtype)
              ),
              running_var_state.num_points,
          ),
          mean=util.map_tree(  # pylint: disable=g-long-lambda
              lambda m, init_m: jnp.ones_like(m) * init_m,
              running_var_state.mean,
              initial_mean,
          ),
          variance=util.map_tree(  # pylint: disable=g-long-lambda
              lambda v, init_s: jnp.ones_like(v) * init_s**2,
              running_var_state.variance,
              initial_scale,
          ),
      )
      ssa_state = step_size_adaptation_init(
          jnp.asarray(step_size, dtype=dtype)
      )
  
      return AdaptiveHamiltonianMonteCarloState(
          hmc_state=hmc_state,
          running_var_state=running_var_state,
          ssa_state=ssa_state,
          step=jnp.zeros([], jnp.int32),
      )
  
imports: import functools
  import typing
  from typing import Any, Callable, Iterable, NamedTuple, Optional
  import numpy
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import malt
  from fun_mc import sga_hmc
  from fun_mc.malt import *
  from fun_mc.sga_hmc import *
label_desc: Use image node_id 4 for calling a global function with example usage: adaptive_hamiltonian_monte_carlo_init(state, target_log_prob_fn, step_size, initial_mean, initial_scale, scale_smoothing_steps) and returns: AdaptiveHamiltonianMonteCarloState