code: def thread_hk_state_in_kwargs(dec_fun):
      """Equivalent to jax.{} but passing Haiku state.""".format(
          dec_fun.__name__
      )
  
      @functools.wraps(dec_fun)
      def wrapped_dec_fun(fun, *dec_args, **dec_kwargs):
          """Decorates a modified version of ``fun`` that passes Haiku state."""
  
          if not base.inside_transform():
              raise ValueError(
                  "hk.{0}() should not be used outside of hk.transform. "
                  "Use jax.{0}() instead.".format(dec_fun.__name__)
              )
  
          @functools.wraps(fun)
          def stateful_fun(*args, **kwargs):
              state_in = kwargs.pop("hk_state")
              with temporary_internal_state(
                  state_in, share_python_state=True
              ), base.push_jax_trace_level():
                  out = fun(*args, **kwargs)
                  return out, difference(state_in, internal_state())
  
          dec_stateful_fun = dec_fun(
              stateful_fun, *dec_args, **dec_kwargs
          )
  
          @functools.wraps(dec_stateful_fun)
          def wrapper(*args, **kwargs):
              kwargs["hk_state"] = internal_state()
              out, state = dec_stateful_fun(*args, **kwargs)
              update_internal_state(state)
              return out
  
          return wrapper
  
      wrapped_dec_fun.__doc__ = f"Equivalent to jax.{dec_fun.__name__} but passing Haiku state."
      return wrapped_dec_fun
  
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 11 for calling a global function with example usage: thread_hk_state_in_kwargs(dec_fun) and returns: wrapped_dec_fun, wrapper, out, out,  difference