code: def step_size_adaptation_init(
      init_step_size: fun_mc.FloatArray,
  ) -> StepSizeAdaptationState:
      """Initializes `StepSizeAdaptationState`.
  
      Args:
        init_step_size: Floating point Array. Initial step size.
  
      Returns:
        step_size_adaptation_state: `StepSizeAdaptationState`
      """
      init_step_size = jnp.asarray(init_step_size)
      rms_state = fun_mc.running_mean_init(
          init_step_size.shape, init_step_size.dtype
      )
      rms_state = rms_state._replace(mean=init_step_size)
  
      return StepSizeAdaptationState(
          step=jnp.array(0, jnp.int32),
          opt_state=fun_mc.adam_init(jnp.log(init_step_size)),
          rms_state=rms_state,
      )
  
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 2 for calling a global function with example usage: step_size_adaptation_init(init_step_size) and returns: StepSizeAdaptationState