code: def push_policy(cls: type[hk.Module], policy: jmp.Policy):
      """Sets the given policy for the given class while the context is active.
  
      Args:
        cls: A Haiku module class.
        policy: A JMP policy to apply to the module.
  
      Yields:
        ``None``.
  
      See also:
        - :func:`clear_policy`: Clears any policies associated with a class.
        - :func:`get_policy`: Gets the policy for a given class.
        - :func:`set_policy`: Sets a policy for a given class.
        - :func:`current_policy`: Retrieves the currently active policy (if any).
      """
      assert policy is not None, "To unset policies use clear_policy."
  
      # Check for trying to push a new policy inside a module method. In theory it
      # is safe to do this when varying the parameter dtype, but we are defensive
      # and ask users to set policies before calling module methods to avoid
      # confusion.
      current_module = base.inside_transform() and base.current_module()
      if current_module and key_for_module(
          type(current_module)
      ) == key_for_module(cls):
          raise ValueError(
              "Pushing a policy inside a method on the same class is not supported."
          )
  
      old_policy = get_policy(cls)
      set_policy(cls, policy)
      try:
          yield
      finally:
          if old_policy is not None:
              set_policy(cls, old_policy)
          else:
              clear_policy(cls)
  
imports: import collections
  import contextlib
  import threading
  from typing import TypeVar, Optional, Union
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  import jmp
label_desc: Use image node_id 6 for calling a global function with example usage: push_policy(cls, policy) without return types