code: def step_size_adaptation_step(
      state: StepSizeAdaptationState,
      log_accept_ratio: fun_mc.FloatArray,
      num_adaptation_steps: Optional[fun_mc.IntArray],
      target_accept_prob: float | fun_mc.FloatArray = 0.8,
      adaptation_rate: float | fun_mc.FloatArray = 0.05,
      adaptation_rate_decay_power: float | fun_mc.FloatArray = 0.1,
      averaging_window_steps: int | fun_mc.IntArray = 100,
      min_log_accept_prob: fun_mc.FloatArray = np.log(1e-5),
      reduce_fn: Callable[[fun_mc.FloatArray], fun_mc.FloatArray] = (
          tfp.math.reduce_logmeanexp
      ),
      opt_kwargs: Optional[dict[str, Any]] = None,
  ) -> tuple[StepSizeAdaptationState, StepSizeAdaptationExtra]:
      """Gradient based step size adaptation using ADAM.
  
      Given the `log_accept_ratio` statistic from an Metropolis-Hastings algorithm,
      this adapts the step size hyperparameter to make that statistic hit some
      `target_accept_prob`. The step size can be extracted using the `step_size`
      method on the state structure.
  
      Args:
        state: `StepSizeAdaptationState`
        log_accept_ratio: Float tensor. The logarithm of the accept ratio.
        num_adaptation_steps: Number of adaptation steps, can be `None`.
        target_accept_prob: Target acceptance probability.
        adaptation_rate: Step size adaptation rate.
        adaptation_rate_decay_power:  Power of the polynomial schedule for
          `trajectory_length_factor` warmup.
        averaging_window_steps: Number of steps to compute the averaged step size.
        min_log_accept_prob: Clamps acceptance probability to this value for
          numerical stability.
        reduce_fn: A function that reduces `log_accept_ratio` in log-space. By
          default, this computes the log-mean-exp.
        opt_kwargs: Additional arguments to pass to the optimizer.
  
      Returns:
        step_size_adaptation_state: `StepSizeAdaptationState`
        step_size_adaptation_extra: `StepSizeAdaptationExtra`
      """
      opt_kwargs = {} if opt_kwargs is None else opt_kwargs
      dtype = log_accept_ratio.dtype
      adaptation_rate = jnp.asarray(adaptation_rate, dtype=dtype)
      target_accept_prob = jnp.asarray(target_accept_prob, dtype=dtype)
      adaptation_rate_decay_power = jnp.asarray(
          adaptation_rate_decay_power, dtype=dtype
      )
      min_log_accept_prob = jnp.full(
          log_accept_ratio.shape, jnp.array(min_log_accept_prob, dtype)
      )
  
      log_accept_prob = jnp.minimum(
          log_accept_ratio, jnp.zeros([], dtype)
      )
      log_accept_prob = jnp.maximum(
          log_accept_prob, min_log_accept_prob
      )
      log_accept_prob = jnp.where(
          jnp.isfinite(log_accept_prob),
          log_accept_prob,
          min_log_accept_prob,
      )
      accept_prob = jnp.exp(reduce_fn(log_accept_prob))
  
      loss_fn = fun_mc.make_surrogate_loss_fn(
          lambda _: (  # pylint: disable=g-long-lambda
              target_accept_prob - accept_prob,
              (),
          )
      )
  
      if num_adaptation_steps is not None:
          adaptation_rate = _polynomial_decay(
              step=state.step,
              step_size=adaptation_rate,
              decay_steps=num_adaptation_steps,
              final_step_size=0.0,
              power=adaptation_rate_decay_power,
          )
  
      # Optimize step size.
      opt_state, opt_extra = fun_mc.adam_step(
          state.opt_state, loss_fn, adaptation_rate, **opt_kwargs
      )
  
      # Do iterate averaging.
      old_rms_state = state.rms_state
      rms_state, _ = fun_mc.running_mean_step(
          old_rms_state,
          jnp.exp(opt_state.state),
          window_size=averaging_window_steps,
      )
  
      if num_adaptation_steps is not None:
          rms_state = util.map_tree(
              lambda n, o: jnp.where(
                  state.step < num_adaptation_steps, n, o
              ),
              rms_state,
              old_rms_state,
          )
  
      state = state._replace(
          opt_state=opt_state, rms_state=rms_state, step=state.step + 1
      )
      extra = StepSizeAdaptationExtra(
          opt_extra=opt_extra, accept_prob=accept_prob
      )
      return state, extra
  
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 3 for calling a global function with example usage: step_size_adaptation_step(state, log_accept_ratio, num_adaptation_steps, target_accept_prob, adaptation_rate, adaptation_rate_decay_power, averaging_window_steps, min_log_accept_prob, reduce_fn, opt_kwargs) and returns: state,  extra