code: def metropolis_adjusted_langevin_trajectories_init(
      state: fun_mc.State, target_log_prob_fn: fun_mc.PotentialFn
  ) -> MetropolisAdjustedLangevinTrajectoriesState:
      """Initializes the `MetropolisAdjustedLangevinTrajectoriesState`.
  
      Args:
        state: State of the chain.
        target_log_prob_fn: Target log prob fn.
  
      Returns:
        malt_state: State of the `metropolis_adjusted_langevin_trajectories_step`
          `TransitionOperator`.
      """
      state = util.map_tree(jnp.array, state)
      target_log_prob, state_extra, state_grads = util.map_tree(
          jnp.array,
          fun_mc.call_potential_fn_with_grads(
              target_log_prob_fn, state
          ),
      )
      return MetropolisAdjustedLangevinTrajectoriesState(
          state=state,
          state_grads=state_grads,
          target_log_prob=target_log_prob,
          state_extra=state_extra,
      )
  
file_path: probability/spinoffs/fun_mc/fun_mc/malt.py
imports: import functools
  from typing import Any, Callable, NamedTuple, Optional, Tuple
  from fun_mc import backend
  from fun_mc import fun_mc_lib
label_desc: Use image node_id 3 for calling a global function with example usage: metropolis_adjusted_langevin_trajectories_init(state, target_log_prob_fn) and returns: MetropolisAdjustedLangevinTrajectoriesState