code: def with_assert_not_used(f):
      """Validates that an updater method is called correctly."""
  
      # NOTE defined outside LiftWithStateUpdater to avoid adding this to the public
      # API and letting users call directly.
      @functools.wraps(f)
      def wrapper(self, *a, **k):
          if self._used:  # pylint: disable=protected-access
              raise ValueError("State updater must only be used once.")
  
          if not base.inside_transform():
              raise ValueError(
                  "State updater must be used inside hk.transform_with_state."
              )
  
          if self._context_id != id(
              base.current_context()
          ):  # pylint: disable=protected-access
              raise ValueError(
                  "State updater must be used within the same call to init/apply."
              )
  
          self._used = True  # pylint: disable=protected-access
  
          return f(self, *a, **k)
  
      return wrapper
  
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 6 for calling a global function with example usage: with_assert_not_used(f) and returns: wrapper, f