code: def add_state_to_init_fn(
      init_fn: Callable[..., hk.Params],
  ) -> Callable[..., tuple[hk.Params, hk.State]]:
      def wrapped_init_fn(*a, **k):
          params = init_fn(*a, **k)
          if not isinstance(params, Mapping):
              raise base.NonEmptyStateError(
                  "For stateful lifted functions use "
                  "`hk.lift_with_state`"
              )
          return params, {}
  
      return wrapped_init_fn
  
imports: from collections.abc import Mapping, MutableMapping
  import functools
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.typing import LiftingModuleType
label_desc: Use image node_id 3 for calling a global function with example usage: add_state_to_init_fn(init_fn) and returns: wrapped_init_fn, params,  dict