code: def update_internal_state(state: InternalState):
      frame = base.current_frame()
      if not frame.params_frozen and state.params is not None:
          update_recursive(frame.params, state.params)
      update_recursive(frame.state, state.state)
      rng = state.rng
      if rng is not None:
          frame.rng_stack.peek().replace_internal_state(rng)
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling a global function with example usage: update_internal_state(state) without return types