code: def adaptive_nuts_step(
      anuts_state: AdaptiveNUTSState,
      target_log_prob_fn: fun_mc.PotentialFn,
      num_mala_steps: int,
      num_adaptation_steps: int,
      seed: jax.Array,
      scalar_step_size: Optional[jnp.ndarray] = None,
      vector_step_size: Optional[jnp.ndarray] = None,
      rvar_factor: int = 8,
      iterate_factor: int = 16,
      target_accept_prob: float = 0.8,
      step_size_adaptation_rate: float = 0.05,
      rvar_smoothing: int = 0,
      step_size_adaptation_rate_decay: str = "none",
      step_size_opt_kwargs: Mapping[
          str, Any
      ] = immutabledict.immutabledict({}),
  ):
      """One step of Adaptive NUTS."""
      seed, mcmc_seed = jax.random.split(seed)
  
      def tfp_target_log_prob_fn(x):
          return target_log_prob_fn(x)[0]
  
      # ===========================
      # Compute common hyper-params
      # ===========================
      adapt = anuts_state.step < num_adaptation_steps
      num_chains = anuts_state.nuts_state[0].shape[0]
  
      if scalar_step_size is None:
          scalar_step_size = jnp.where(
              adapt,
              jnp.exp(anuts_state.log_step_size_opt_state.state),
              anuts_state.step_size_rmean_state.mean,
          )
      if vector_step_size is None:
          vector_step_size = jnp.sqrt(anuts_state.rvar_state.variance)
          vector_step_size = vector_step_size / vector_step_size.max()
  
      trans_kernel = tfp.mcmc.NoUTurnSampler(
          target_log_prob_fn=tfp_target_log_prob_fn,
          step_size=scalar_step_size * vector_step_size,
      )
      trans_kernel._max_tree_depth = jnp.where(
          anuts_state.step < num_mala_steps,
          1,  # pylint: disable=protected-access
          10,
      )
  
      # ====================
      # Apply a step of MCMC
      # ====================
      state, kr = anuts_state.nuts_state
      kr = kr._replace(step_size=scalar_step_size * vector_step_size)
      nuts_state = trans_kernel.one_step(state, kr, seed=mcmc_seed)
  
      # ==========
      # Adaptation
      # ==========
  
      # Adjust running-variance estimate.
      cand_rvar_state, _ = fun_mc.running_variance_step(
          anuts_state.rvar_state,
          nuts_state[0],
          axis=0,
          window_size=rvar_smoothing
          + num_chains * anuts_state.step // rvar_factor,
      )
      rvar_state = fun_mc.choose(
          adapt, cand_rvar_state, anuts_state.rvar_state
      )
  
      # Adjust step size.
      log_accept_ratio = nuts_state[1].log_accept_ratio
      min_log_accept_prob = jnp.full(
          log_accept_ratio.shape, jnp.log(1e-5)
      )
  
      log_accept_prob = jnp.minimum(log_accept_ratio, jnp.zeros([]))
      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(
          tfp.math.reduce_log_harmonic_mean_exp(log_accept_prob)
      )
  
      log_step_size_surrogate_loss_fn = fun_mc.make_surrogate_loss_fn(
          lambda _: (target_accept_prob - accept_prob, ())
      )
  
      cur_adaptation_rate = jnp.where(
          adapt, step_size_adaptation_rate, 0.0
      )
      if step_size_adaptation_rate_decay == "linear":
          cur_adaptation_rate = (
              1.0 - 0.95 * anuts_state.step / num_adaptation_steps
          ) * cur_adaptation_rate
      (log_step_size_opt_state, _) = fun_mc.adam_step(
          anuts_state.log_step_size_opt_state,
          log_step_size_surrogate_loss_fn,
          cur_adaptation_rate,
          **step_size_opt_kwargs,
      )
  
      # =================
      # Iterate averaging
      # =================
      cand_step_size_rmean_state, _ = fun_mc.running_mean_step(
          anuts_state.step_size_rmean_state,
          scalar_step_size,
          window_size=anuts_state.step // iterate_factor,
      )
      step_size_rmean_state = fun_mc.choose(
          adapt,
          cand_step_size_rmean_state,
          anuts_state.step_size_rmean_state,
      )
  
      anuts_state = anuts_state._replace(
          nuts_state=nuts_state,
          rvar_state=rvar_state,
          log_step_size_opt_state=log_step_size_opt_state,
          step_size_rmean_state=step_size_rmean_state,
          step=anuts_state.step + 1,
      )
      anuts_extra = AdaptiveNUTSExtra(
          scalar_step_size=scalar_step_size,
          vector_step_size=vector_step_size,
          num_integrator_steps=nuts_state[1].leapfrogs_taken.max(),
          mean_num_integrator_steps=nuts_state[
              1
          ].leapfrogs_taken.mean(),
          log_accept_ratio=log_accept_ratio,
          extra={},
      )
      return anuts_state, anuts_extra
  
imports: import functools
  import os
  from typing import Any, Callable, Dict, Optional, Mapping, NamedTuple, Tuple, Union
  from etils import epath
  import gin
  import immutabledict
  import jax
  import jax.numpy
  import numpy
  from discussion.adaptive_malt import utils
  from fun_mc import using_jax
  from inference_gym import using_jax
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 9 for calling a global function with example usage: adaptive_nuts_step(anuts_state, target_log_prob_fn, num_mala_steps, num_adaptation_steps, seed, scalar_step_size, vector_step_size, rvar_factor, iterate_factor, target_accept_prob, step_size_adaptation_rate, rvar_smoothing, step_size_adaptation_rate_decay, step_size_opt_kwargs) and returns: anuts_state,  anuts_extra, target_log_prob_fn