code: def adaptive_mcmc_step(
      amcmc_state: AdaptiveMCMCState,
      target_log_prob_fn: fun_mc.PotentialFn,
      num_mala_steps: int,
      num_adaptation_steps: int,
      seed: jax.Array,
      method: str = "hmc",
      damping: Optional[jnp.ndarray] = None,
      scalar_step_size: Optional[jnp.ndarray] = None,
      vector_step_size: Optional[jnp.ndarray] = None,
      mean_trajectory_length: Optional[jnp.ndarray] = None,
      power: Optional[jnp.ndarray] = None,
      principal: Optional[jnp.ndarray] = None,
      max_num_integrator_steps: int = 500,
      rvar_factor: int = 8,
      iterate_factor: int = 16,
      principal_factor: int = 8,
      target_accept_prob: float = 0.8,
      step_size_adaptation_rate: float = 0.05,
      trajectory_length_adaptation_rate: float = 0.05,
      damping_factor: float = 1.0,
      principal_mean_method: str = "running_mean",
      state_grad_estimator: str = "two_dir",
      adapt_normalization_power: bool = False,
      mix_state_mean_in_snaper: bool = False,
      step_size_adaptation_rate_decay: str = "none",
      trajectory_length_adaptation_rate_decay: str = "none",
      rvar_smoothing: int = 0,
      jitter_style: str = "halton",
      use_precond_eigenvalue: bool = False,
      trajectory_opt_kwargs: Mapping[
          str, Any
      ] = immutabledict.immutabledict({}),
      step_size_opt_kwargs: Mapping[
          str, Any
      ] = immutabledict.immutabledict({}),
  ):
      """Applies a step of the adaptive MCMC algorithm.
  
      Args:
        amcmc_state: Adaptive MCMC state.
        target_log_prob_fn: Target log prob function.
        num_mala_steps: Number of initial steps to be MALA.
        num_adaptation_steps: Number of steps to adapt hyperparameters.
        seed: Random seed.
        method: MCMC method. Either 'hmc', 'malt'.
        damping: If not None, the fixed damping to use.
        scalar_step_size: If not None, the fixed scalar step size to use.
        vector_step_size: If not None, the fixed vector step size to use.
        mean_trajectory_length: If not None, the fixed mean trajectory length to
          use.
        power: Power used in the SNAPER criterion.
        principal: If not None, the fixed unnormalized principal component to use.
        max_num_integrator_steps: Maximum number of integrator steps.
        rvar_factor: Factor for running variance adaptation rate.
        iterate_factor: Factor for iterate averaging rate.
        principal_factor: Factor for principal component adaptation rate.
        target_accept_prob: Target acceptance probability.
        step_size_adaptation_rate: Scalar step size adaptation rate.
        trajectory_length_adaptation_rate: Trajectory length adaptation rate.
        damping_factor: Factor for computing the damping coefficient.
        principal_mean_method: Method for computing principal mean. Either
          'running_mean' or 'state_mean'. the running mean for preconditioning:
        state_grad_estimator: State grad estimator to use. Can be 'one_dir' or
          'two_dir'.
        adapt_normalization_power: Whether to adapt the power used for trajectory
          length normalization term in the snaper criterion.
        mix_state_mean_in_snaper: Whether to mix in the current state mean when
          computing the criterion.
        step_size_adaptation_rate_decay: How to decay the adaptation rate. One of
          'none' or 'linear'.
        trajectory_length_adaptation_rate_decay: How to decay the adaptation rate.
          One of 'none' or 'linear'.
        rvar_smoothing: Smoothing points.
        jitter_style: HMC jitter style. One of 'halton' or 'exponential'.
        use_precond_eigenvalue: Whether to use the preconditioned eigenvalue when
          computing damping.
        trajectory_opt_kwargs: Extra arguments to the trajectory length optimizer.
        step_size_opt_kwargs: Extra arguments to the step size optimizer.
  
      Returns:
        A tuple of new adaptive MCMC state and extra.
      """
      seed, mcmc_seed = jax.random.split(seed)
  
      # ===========================
      # Compute common hyper-params
      # ===========================
      adapt = amcmc_state.step < num_adaptation_steps
      mala = amcmc_state.step < num_mala_steps
      num_chains = amcmc_state.mcmc_state.state.shape[0]
  
      if scalar_step_size is None:
          scalar_step_size = jnp.where(
              adapt,
              jnp.exp(amcmc_state.log_step_size_opt_state.state),
              amcmc_state.step_size_rmean_state.mean,
          )
      if vector_step_size is None:
          vector_step_size = jnp.sqrt(amcmc_state.rvar_state.variance)
          vector_step_size = vector_step_size / vector_step_size.max()
  
      if power is None:
          if adapt_normalization_power:
              var = 0.5 * (
                  amcmc_state.proj_rautocov_state.covariance[0, 0]
                  + amcmc_state.proj_rautocov_state.covariance[1, 1]
              )
              cov = amcmc_state.proj_rautocov_state.covariance[1, 0]
              power = 0.5 * (1.0 + cov / var)
          else:
              power = 1.0
  
      if principal is None:
          max_eigenvalue = jnp.linalg.norm(
              amcmc_state.principal_rmean_state.mean
          )
          principal = (
              amcmc_state.principal_rmean_state.mean / max_eigenvalue
          )
      else:
          max_eigenvalue = jnp.linalg.norm(principal)
          principal = principal / max_eigenvalue
  
      if use_precond_eigenvalue:
          max_eigenvalue = jnp.linalg.norm(
              amcmc_state.precond_principal_rmean_state.mean
          )
  
      if method == "hmc":
          if jitter_style == "halton":
              traj_factor = (
                  halton(amcmc_state.step.astype(jnp.float32)) * 2.0
              )
          elif jitter_style == "exponential":
              traj_factor = tfd.Exponential(1.0).sample(
                  seed=jax.random.PRNGKey(amcmc_state.step)
              )  # pytype: disable=wrong-arg-types  # jax-ndarray
          elif jitter_style == "halton_exponential":
              traj_factor = -jnp.log(
                  halton(amcmc_state.step.astype(jnp.float32))
              )
      elif method == "malt":
          traj_factor = 1.0
  
      if mean_trajectory_length is None:
          mean_trajectory_length = jnp.where(
              adapt,
              jnp.exp(
                  amcmc_state.log_trajectory_length_opt_state.state
              ),
              amcmc_state.trajectory_length_rmean_state.mean,
          )
      trajectory_length = traj_factor * mean_trajectory_length
      num_integrator_steps = jnp.ceil(
          trajectory_length / scalar_step_size
      )
      num_integrator_steps = jnp.where(
          jnp.isfinite(num_integrator_steps), num_integrator_steps, 1
      )
      num_integrator_steps = num_integrator_steps.astype(jnp.int32)
  
      # ====================
      # Apply a step of MCMC
      # ====================
      if method == "hmc":
          mcmc_state, mcmc_extra = fun_mc.hamiltonian_monte_carlo_step(
              amcmc_state.mcmc_state,
              target_log_prob_fn=target_log_prob_fn,
              step_size=scalar_step_size * vector_step_size,
              num_integrator_steps=num_integrator_steps,
              seed=mcmc_seed,
          )
          proposed_state = mcmc_extra.proposed_hmc_state.state
      elif method == "malt":
          if damping is None:
              damping = damping_factor / (
                  1e-10 + jnp.sqrt(max_eigenvalue)
              )
          (
              mcmc_state,
              mcmc_extra,
          ) = fun_mc.prefab.metropolis_adjusted_langevin_trajectories_step(
              amcmc_state.mcmc_state,
              target_log_prob_fn=target_log_prob_fn,
              step_size=scalar_step_size * vector_step_size,
              momentum_refresh_fn=lambda m, seed: (  # pylint: disable=g-long-lambda
                  malt_lib._gaussian_momentum_refresh_fn(  # pylint: disable=protected-access
                      m,
                      damping=damping,
                      step_size=scalar_step_size / 2.0,
                      seed=seed,
                  )
              ),
              num_integrator_steps=num_integrator_steps,
              damping=damping,
              seed=mcmc_seed,
          )
          proposed_state = mcmc_extra.proposed_malt_state.state
      # This assumes a N(0, I) momentum distribution
      initial_momentum = mcmc_extra.initial_momentum
      final_momentum = mcmc_extra.integrator_extra.momentum_grads
  
      # ==========
      # Adaptation
      # ==========
  
      # Adjust running-variance estimate.
      cand_rvar_state, _ = fun_mc.running_variance_step(
          amcmc_state.rvar_state,
          mcmc_state.state,
          axis=0,
          window_size=rvar_smoothing
          + num_chains * amcmc_state.step // rvar_factor,
      )
      rvar_state = fun_mc.choose(
          adapt, cand_rvar_state, amcmc_state.rvar_state
      )
  
      # Adjust trajectory length.
      def log_trajectory_length_surrogate_loss_fn(
          log_trajectory_length,
      ):
          log_trajectory_length = fun_mc.clip_grads(
              log_trajectory_length, 1e6
          )
          mean_trajectory_length = jnp.exp(log_trajectory_length)
          trajectory_length = traj_factor * mean_trajectory_length
  
          if state_grad_estimator == "one_dir":
              start_end_grads = [
                  (
                      amcmc_state.mcmc_state.state,
                      proposed_state,
                      final_momentum,
                  ),
              ]
          elif state_grad_estimator == "two_dir":
              start_end_grads = [
                  (
                      amcmc_state.mcmc_state.state,
                      proposed_state,
                      final_momentum,
                  ),
                  (
                      proposed_state,
                      amcmc_state.mcmc_state.state,
                      -initial_momentum,
                  ),
              ]
  
          criteria = []
          extras = []
          for start_state, end_state, state_grads in start_end_grads:
              action = (
                  trajectory_length * vector_step_size * state_grads
              )
              end_state_with_action = (
                  end_state + action - jax.lax.stop_gradient(action)
              )
  
              if mix_state_mean_in_snaper:
                  state_mean_weight = (
                      amcmc_state.rvar_state.num_points
                  ) / (amcmc_state.rvar_state.num_points + num_chains)
              else:
                  state_mean_weight = 1.0
  
              criterion, extra = snaper_criterion(
                  previous_state=start_state,
                  proposed_state=end_state_with_action,
                  accept_prob=jnp.exp(
                      jnp.minimum(0.0, -mcmc_extra.log_accept_ratio)
                  ),
                  trajectory_length=trajectory_length
                  + scalar_step_size,
                  principal=principal,
                  power=power,
                  # These two expressions are a bit weird for the reverse direction...
                  state_mean=amcmc_state.rvar_state.mean,
                  state_mean_weight=state_mean_weight,
              )
              criteria.append(criterion)
              extras.append(extra)
          return -jnp.array(criteria).mean(), extras
  
      cur_adaptation_rate = jnp.where(
          adapt, trajectory_length_adaptation_rate, 0.0
      )
      if trajectory_length_adaptation_rate_decay == "linear":
          cur_adaptation_rate = (
              1.0 - 0.95 * amcmc_state.step / num_adaptation_steps
          ) * cur_adaptation_rate
      (
          log_trajectory_length_opt_state,
          log_trajectory_length_opt_extra,
      ) = fun_mc.adam_step(
          amcmc_state.log_trajectory_length_opt_state,
          log_trajectory_length_surrogate_loss_fn,
          cur_adaptation_rate,
          **trajectory_opt_kwargs,
      )
      log_trajectory_length_opt_state = (
          log_trajectory_length_opt_state._replace(
              state=jnp.where(
                  mala,
                  jnp.log(scalar_step_size),
                  log_trajectory_length_opt_state.state,
              )
          )
      )
  
      # Adjust step size.
      log_accept_ratio = mcmc_extra.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 * amcmc_state.step / num_adaptation_steps
          ) * cur_adaptation_rate
      (log_step_size_opt_state, _) = fun_mc.adam_step(
          amcmc_state.log_step_size_opt_state,
          log_step_size_surrogate_loss_fn,
          cur_adaptation_rate,
          **step_size_opt_kwargs,
      )
  
      # Clip the trajectory length (need new step size estimate)
      log_trajectory_length_opt_state = (
          log_trajectory_length_opt_state._replace(
              state=jnp.minimum(
                  jnp.log(scalar_step_size * max_num_integrator_steps),
                  log_trajectory_length_opt_state.state,
              )
          )
      )
  
      # Adjust the principal component + max_eigval estimate.
      if principal_mean_method == "state_mean":
          mean = mcmc_state.state.mean(0)
      elif principal_mean_method == "running_mean":
          mean = rvar_state.mean
  
      cand_principal_rmean_state, _ = fun_mc.running_mean_step(
          amcmc_state.principal_rmean_state,
          ccipca(
              mcmc_state.state - mean,
              amcmc_state.principal_rmean_state.mean,
          ),
          window_size=rvar_smoothing
          + amcmc_state.step // principal_factor,
      )
      principal_rmean_state = fun_mc.choose(
          adapt,
          cand_principal_rmean_state,
          amcmc_state.principal_rmean_state,
      )
  
      cand_precond_principal_rmean_state, _ = fun_mc.running_mean_step(
          amcmc_state.precond_principal_rmean_state,
          ccipca(
              (mcmc_state.state - mean) / vector_step_size,
              amcmc_state.precond_principal_rmean_state.mean,
          ),
          window_size=rvar_smoothing
          + amcmc_state.step // principal_factor,
      )
      precond_principal_rmean_state = fun_mc.choose(
          adapt,
          cand_precond_principal_rmean_state,
          amcmc_state.precond_principal_rmean_state,
      )
  
      # Adjust auto-covariance of the squared projections.
      cand_proj_rautocov_state, _ = fun_mc.running_covariance_step(
          amcmc_state.proj_rautocov_state,
          jnp.stack(
              [
                  log_trajectory_length_opt_extra.loss_extra[0][
                      "proposed_projection"
                  ]
                  ** 2,
                  log_trajectory_length_opt_extra.loss_extra[0][
                      "previous_projection"
                  ]
                  ** 2,
              ],
              -1,
          ),
          axis=0,
          window_size=rvar_smoothing
          + num_chains * amcmc_state.step // rvar_factor,
      )
      proj_rautocov_state = fun_mc.choose(
          adapt,
          cand_proj_rautocov_state,
          amcmc_state.proj_rautocov_state,
      )
  
      # =================
      # Iterate averaging
      # =================
      cand_step_size_rmean_state, _ = fun_mc.running_mean_step(
          amcmc_state.step_size_rmean_state,
          scalar_step_size,
          window_size=amcmc_state.step // iterate_factor,
      )
      step_size_rmean_state = fun_mc.choose(
          adapt,
          cand_step_size_rmean_state,
          amcmc_state.step_size_rmean_state,
      )
  
      cand_trajectory_length_rmean_state, _ = fun_mc.running_mean_step(
          amcmc_state.trajectory_length_rmean_state,
          mean_trajectory_length,
          window_size=amcmc_state.step // iterate_factor,
      )
      trajectory_length_rmean_state = fun_mc.choose(
          adapt,
          cand_trajectory_length_rmean_state,
          amcmc_state.trajectory_length_rmean_state,
      )
  
      proposed_projection2 = ((proposed_state - mean) * principal).sum(
          -1
      )
      previous_projection2 = (
          (amcmc_state.mcmc_state.state - mean) * principal
      ).sum(-1)
      amcmc_state = amcmc_state._replace(
          mcmc_state=mcmc_state,
          rvar_state=rvar_state,
          proj_rautocov_state=proj_rautocov_state,
          principal_rmean_state=principal_rmean_state,
          precond_principal_rmean_state=precond_principal_rmean_state,
          log_step_size_opt_state=log_step_size_opt_state,
          log_trajectory_length_opt_state=log_trajectory_length_opt_state,
          step_size_rmean_state=step_size_rmean_state,
          trajectory_length_rmean_state=trajectory_length_rmean_state,
          step=amcmc_state.step + 1,
      )
      amcmc_extra = AdaptiveMCMCExtra(
          mcmc_extra=mcmc_extra,
          scalar_step_size=scalar_step_size,
          vector_step_size=vector_step_size,
          principal=principal,
          power=power,
          max_eigenvalue=max_eigenvalue,
          damping=damping,
          mean_trajectory_length=mean_trajectory_length,
          log_trajectory_length_opt_extra=log_trajectory_length_opt_extra,
          num_integrator_steps=num_integrator_steps,
          extra={
              "rautocov": proj_rautocov_state.covariance,
              "proposed_projection": log_trajectory_length_opt_extra.loss_extra[
                  0
              ][
                  "proposed_projection"
              ],
              "previous_projection": log_trajectory_length_opt_extra.loss_extra[
                  0
              ][
                  "previous_projection"
              ],
              "proposed_projection2": proposed_projection2,
              "previous_projection2": previous_projection2,
              "mean": mean,
              "principal_rmean": principal_rmean_state.mean,
          },
      )
      return amcmc_state, amcmc_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 7 for calling a global function with example usage: adaptive_mcmc_step(amcmc_state, target_log_prob_fn, num_mala_steps, num_adaptation_steps, seed, method, damping, scalar_step_size, vector_step_size, mean_trajectory_length, power, principal, max_num_integrator_steps, rvar_factor, iterate_factor, principal_factor, target_accept_prob, step_size_adaptation_rate, trajectory_length_adaptation_rate, damping_factor, principal_mean_method, state_grad_estimator, adapt_normalization_power, mix_state_mean_in_snaper, step_size_adaptation_rate_decay, trajectory_length_adaptation_rate_decay, rvar_smoothing, jitter_style, use_precond_eigenvalue, trajectory_opt_kwargs, step_size_opt_kwargs) and returns: amcmc_state,  amcmc_extra, unknown,  extras