code: def meads_step(
      meads_state: MeadsState,
      target_log_prob_fn: fun_mc.PotentialFn,
      seed: jax.Array,
      vector_step_size: Optional[jnp.ndarray] = None,
      damping: Optional[jnp.ndarray] = None,
      step_size_multiplier: float = 0.5,
      damping_slowdown: float = 1.0,
  ):
      """One step of MEADS."""
      phmc_seed, seed = jax.random.split(seed)
  
      num_folds = meads_state.phmc_state.state.shape[0]
      chains_per_fold = meads_state.phmc_state.state.shape[1]
      num_chains = num_folds * chains_per_fold
      step = meads_state.step
      fold_to_skip = meads_state.fold_to_skip
      phmc_state = meads_state.phmc_state
  
      # Randomly refold the walkers.
      perm = jax.random.permutation(
          jax.random.PRNGKey(step // 4), num_chains
      )  # pytype: disable=wrong-arg-types  # jax-ndarray
      # TODO(mhoffman): This should really done with a scatter.
      unperm = jnp.eye(num_chains)[perm].argmax(0)
  
      def refold(x, perm):
          return x.reshape((num_chains,) + x.shape[2:])[perm].reshape(
              x.shape
          )
  
      phmc_state = jax.tree_map(
          functools.partial(refold, perm=perm), phmc_state
      )
  
      if vector_step_size is None:
          vector_step_size = phmc_state.state.std(1, keepdims=True)
      else:
          vector_step_size = jnp.tile(
              vector_step_size[jnp.newaxis, jnp.newaxis],
              [num_folds, 1, 1],
          )
      # Apply preconditioning within-fold to estimate step size and damping.
      self_preconditioned_state = phmc_state.state / vector_step_size
      self_preconditioned_grads = (
          phmc_state.state_grads * vector_step_size
      )
      # Apply preconditioning across folds to do the actual update.
      rolled_vector_step_size = jnp.roll(vector_step_size, 1, 0)
  
      # Set step size for each fold based on the fold to its left.
      scalar_step_size = step_size_multiplier / jnp.sqrt(
          jax.vmap(
              functools.partial(
                  estimate_largest_eigenvalue_of_covariance,
                  remove_mean=False,
              )
          )(self_preconditioned_grads)
      )
      scalar_step_size = jnp.minimum(1.0, scalar_step_size)
      scalar_step_size = jnp.roll(scalar_step_size, 1)
  
      # Set damping.
      max_eigenvalue = jax.vmap(
          estimate_largest_eigenvalue_of_covariance
      )(self_preconditioned_state)
      if damping is None:
          damping = scalar_step_size / jnp.sqrt(max_eigenvalue)
      else:
          damping = jnp.broadcast_to(damping, num_folds)
  
      # Put a floor on the amount of damping in early iterations.
      damping = jnp.maximum(damping_slowdown / step, damping)
  
      noise_fraction = (1 - jnp.exp(-2 * damping)) ** 0.5
      mh_drift = 0.5 * noise_fraction**2
  
      # TODO(mhoffman): Consider using lax.gather instead of jnp.roll logic.
      # An advantage of roll is that it makes it clear to XLA that there's not
      # actually any dynamic sizing going on here.
      def select_folds(x):
          return jnp.roll(
              jnp.roll(x, -fold_to_skip, 0)[1:], fold_to_skip, 0
          )
  
      def rejoin_folds(updated, original):
          return jnp.roll(
              jnp.concatenate(
                  [
                      original[fold_to_skip][jnp.newaxis],
                      jnp.roll(updated, -fold_to_skip, 0),
                  ],
                  0,
              ),
              fold_to_skip,
              0,
          )
  
      (
          active_fold_state,
          phmc_extra,
      ) = fun_mc.prefab.persistent_hamiltonian_monte_carlo_step(
          jax.tree_map(select_folds, phmc_state),
          target_log_prob_fn=target_log_prob_fn,
          step_size=select_folds(
              scalar_step_size[:, jnp.newaxis, jnp.newaxis]
              * rolled_vector_step_size
          ),
          num_integrator_steps=1,
          noise_fraction=select_folds(noise_fraction)[
              :, jnp.newaxis, jnp.newaxis
          ],
          mh_drift=select_folds(mh_drift)[:, jnp.newaxis],
          seed=phmc_seed,
      )
      phmc_state = jax.tree_map(
          rejoin_folds, active_fold_state, phmc_state
      )
  
      # Revert the ordering of the walkers.
      phmc_state = jax.tree_map(
          functools.partial(refold, perm=unperm), phmc_state
      )
  
      meads_state = MeadsState(
          phmc_state=phmc_state,
          fold_to_skip=(fold_to_skip + 1) % num_folds,
          step=step + 1,
      )
  
      extra = (
          MeadsExtra(  # pytype: disable=wrong-arg-types  # jax-ndarray
              phmc_extra=phmc_extra,
              scalar_step_size=scalar_step_size,
              vector_step_size=vector_step_size,
              damping=damping,
              max_eigenvalue=max_eigenvalue,
              num_integrator_steps=1,
              level=phmc_state.pmh_state.level,
              extra={},
          )
      )
  
      return meads_state, 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 15 for calling a global function with example usage: meads_step(meads_state, target_log_prob_fn, seed, vector_step_size, damping, step_size_multiplier, damping_slowdown) and returns: meads_state,  extra, x, jnp, jnp