code: def stochastic_gradient_ascent_hmc_init(
      state: fun_mc.State,
      target_log_prob_fn: fun_mc.PotentialFn,
      init_trajectory_length: fun_mc.FloatArray,
      trajectory_length_params_init_fn: Callable[
          [fun_mc.FloatArray], Any
      ] = default_trajectory_length_init,
  ):
      """Initialize Stochastic Gradient Ascent HMC state.
  
      Args:
        state: Initial Markov Chain state.
        target_log_prob_fn: Target log prob fn.
        init_trajectory_length: Initial trajectory length. Passed to
          `trajectory_length_params_init_fn`.
        trajectory_length_params_init_fn: Initializer for the trajectory length
          parameters.
  
      Returns:
        sga_hmc_state: New Stochastic Gradient Ascent HMC state.
      """
      init_trajectory_length = jnp.asarray(init_trajectory_length)
      init_trajectory_length_params = trajectory_length_params_init_fn(
          init_trajectory_length
      )
      return StochasticGradientAscentHMCState(
          hmc_state=fun_mc.hamiltonian_monte_carlo_init(
              state, target_log_prob_fn
          ),
          step=jnp.ones([], jnp.int32),
          trajectory_length_params_opt_state=fun_mc.adam_init(
              init_trajectory_length_params
          ),
          trajectory_length_params_rmean_state=fun_mc.running_mean_init(
              util.map_tree(
                  lambda x: x.shape, init_trajectory_length_params
              ),
              util.map_tree(
                  lambda x: x.dtype, init_trajectory_length_params
              ),
          )._replace(mean=init_trajectory_length_params),
      )
  
imports: from collections.abc import Mapping
  from typing import Any, Callable, NamedTuple, Optional
  import immutabledict
  from fun_mc import backend
  from fun_mc import fun_mc_lib
label_desc: Use image node_id 8 for calling a global function with example usage: stochastic_gradient_ascent_hmc_init(state, target_log_prob_fn, init_trajectory_length, trajectory_length_params_init_fn) and returns: StochasticGradientAscentHMCState