code: def set_policy(cls: type[hk.Module], policy: jmp.Policy):
      """Uses the given policy for all instances of the module class.
  
      NOTE: Policies are only applied to modules created in the current thread.
  
      A mixed precision policy describes how inputs, module parameters and module
      outputs should be cast at runtime. By applying a policy to a given type of
      module, you can control how all instances of that module behave in your
      program.
  
      For example, you might want to try running a ResNet50 model in a mixture of
      ``float16`` and ``float32`` on GPU to get higher throughput. To do so you can
      apply a mixed precision policy to the ResNet50 type that will create
      parameters in ``float32``, but cast them to ``float16`` before use, along with
      all module inputs:
  
      >>> policy = jmp.get_policy('params=float32,compute=float16,output=float32')
      >>> hk.mixed_precision.set_policy(hk.nets.ResNet50, policy)
      >>> net = hk.nets.ResNet50(4)
      >>> x = jnp.ones([4, 224, 224, 3])
      >>> print(net(x, is_training=True))
      [[0. 0. 0. 0.]
       [0. 0. 0. 0.]
       [0. 0. 0. 0.]
       [0. 0. 0. 0.]]
  
      For a fully worked mixed precision example see the imagenet example in Haiku's
      examples directory. This example shows mixed precision on GPU offering a 2x
      speedup in training time with only a small impact on final top-1 accuracy.
  
      >>> hk.mixed_precision.clear_policy(hk.nets.ResNet50)
  
      Args:
        cls: A Haiku module class.
        policy: A JMP policy to apply to the module.
  
      See also:
        - :func:`push_policy`: Context manager for setting policies.
        - :func:`current_policy`: Retrieves the currently active policy (if any).
        - :func:`clear_policy`: Clears any policies associated with a class.
        - :func:`get_policy`: Gets the policy for a given class.
      """
      assert policy is not None, "To unset policies use clear_policy."
      _thread_local_state.set_policy(cls, policy)
  
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 5 for calling a global function with example usage: set_policy(cls, policy) without return types