code: def sig_add_state(s: inspect.Signature) -> inspect.Signature:
      """Add hk.State to the return type of a signature."""
      if s.return_annotation is inspect.Parameter.empty:
          ret = Any
      else:
          ret = s.return_annotation
      return inspect.Signature(
          parameters=list(s.parameters.values()),
          return_annotation=tuple[ret, hk.State],
          __validate_parameters__=False,
      )
  
imports: from collections.abc import Mapping
  import inspect
  from typing import Any, Callable, NamedTuple, Optional, TypeVar, Union
  from haiku._src import analytics
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import typing
  import jax
label_desc: Use image node_id 3 for calling a global function with example usage: sig_add_state(s) and returns: inspect