code: def make_unrollable_network(
      make_core_module: Callable[[], hk.RNNCore] = hk.IdentityCore,
      make_feedforward_module: Optional[
          Callable[[], hk.SupportsCall]
      ] = None,
      make_initial_state_fn: Optional[
          Callable[[], hk.SupportsCall]
      ] = None,
  ) -> UnrollableNetwork:
      """Produces an UnrollableNetwork and a state initializing hk.Transformed."""
  
      def default_initial_state_fn(
          batch_size: Optional[int] = None,
      ) -> jnp.ndarray:
          return make_core_module().initial_state(batch_size)
  
      def _apply_core_fn(
          observation: types.Observation, state: jnp.ndarray
      ) -> Tuple[jnp.ndarray, jnp.ndarray]:
          if make_feedforward_module:
              observation = make_feedforward_module()(observation)
          return make_core_module()(observation, state)
  
      def _unroll_core_fn(
          observation: types.Observation, state: jnp.ndarray
      ) -> Tuple[jnp.ndarray, jnp.ndarray]:
          if make_feedforward_module:
              observation = make_feedforward_module()(observation)
          return hk.dynamic_unroll(
              make_core_module(), observation, state
          )
  
      if make_initial_state_fn:
          initial_state_fn = make_initial_state_fn()
      else:
          initial_state_fn = default_initial_state_fn
  
      # Transform module functions into pure functions.
      hk_initial_state_fn = hk.without_apply_rng(
          hk.transform(initial_state_fn)
      )
      apply_core = hk.without_apply_rng(hk.transform(_apply_core_fn))
      unroll_core = hk.without_apply_rng(hk.transform(_unroll_core_fn))
  
      # Pack all core network pure functions into a single convenient container.
      return UnrollableNetwork(
          init=apply_core.init,
          apply=apply_core.apply,
          unroll=unroll_core.apply,
          initial_state_fn_init=hk_initial_state_fn.init,
          initial_state_fn=hk_initial_state_fn.apply,
      )
  
imports: import dataclasses
  from typing import Callable, NamedTuple, Optional, Sequence, Tuple, Union
  from acme import specs
  from acme.agents.jax.mpo import types
  from acme.jax import networks
  from acme.jax import utils
  import chex
  import haiku
  import haiku.initializers
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 2 for calling a global function with example usage: make_unrollable_network(make_core_module, make_feedforward_module, make_initial_state_fn) and returns: UnrollableNetwork, make_core_module, make_core_module, hk