code: def make_unrollable_network(
      environment_spec: specs.EnvironmentSpec,
      make_core_module: Callable[[], hk.RNNCore],
  ) -> UnrollableNetwork:
      """Builds an UnrollableNetwork from a hk.Module factory."""
  
      dummy_observation = jax_utils.zeros_like(
          environment_spec.observations
      )
  
      def make_unrollable_network_functions():
          model = make_core_module()
          apply = model.__call__
  
          def init() -> Tuple[NetworkOutput, RecurrentState]:
              return model(dummy_observation, model.initial_state(None))
  
          return init, (
              apply,
              model.unroll,
              model.initial_state,
          )  # pytype: disable=attribute-error
  
      # Transform and unpack pure functions
      f = hk.multi_transform(make_unrollable_network_functions)
      apply, unroll, initial_state_fn = f.apply
  
      def init_recurrent_state(
          key: jax_types.PRNGKey, batch_size: Optional[int]
      ) -> RecurrentState:
          # TODO(b/244311990): Consider supporting parameterized and learnable initial
          # state functions.
          no_params = None
          return initial_state_fn(no_params, key, batch_size)
  
      return UnrollableNetwork(
          f.init, apply, unroll, init_recurrent_state
      )
  
file_path: acme/acme/jax/networks/base.py
imports: import dataclasses
  from typing import Callable, Optional, Tuple
  from acme import specs
  from acme import types
  from acme.jax import types
  from acme.jax import utils
  import haiku
  import jax.numpy
  from typing_extensions import Protocol
label_desc: Use image node_id 2 for calling a global function with example usage: make_unrollable_network(environment_spec, make_core_module) and returns: UnrollableNetwork, init,  apply,  model,  model, initial_state_fn, model