code: def persistent_hamiltonian_monte_carlo_init(
      state: fun_mc.ArrayNest,
      target_log_prob_fn: fun_mc.PotentialFn,
      momentum: Optional[fun_mc.State] = None,
      init_level: float | fun_mc.FloatArray = 0.0,
  ) -> PersistentHamiltonianMonteCarloState:
      """Initializes the `PersistentHamiltonianMonteCarloState`.
  
      Args:
        state: State of the chain.
        target_log_prob_fn: Target log prob fn.
        momentum: Initial momentum. Set to all zeros by default.
        init_level: Initial level for the persistent Metropolis Hastings.
  
      Returns:
        hmc_state: `PersistentMetropolistHastingsState`.
      """
      state = util.map_tree(jnp.asarray, state)
      target_log_prob, state_extra, state_grads = util.map_tree(
          jnp.asarray,
          fun_mc.call_potential_fn_with_grads(
              target_log_prob_fn, state
          ),
      )
      return PersistentHamiltonianMonteCarloState(
          state=state,
          state_grads=state_grads,
          momentum=momentum
          if momentum is not None
          else util.map_tree(jnp.zeros_like, state),
          target_log_prob=target_log_prob,
          state_extra=state_extra,
          direction=jnp.ones_like(target_log_prob),
          pmh_state=fun_mc.persistent_metropolis_hastings_init(
              shape=target_log_prob.shape,
              dtype=target_log_prob.dtype,
              init_level=init_level,
          ),
      )
  
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 9 for calling a global function with example usage: persistent_hamiltonian_monte_carlo_init(state, target_log_prob_fn, momentum, init_level) and returns: PersistentHamiltonianMonteCarloState