code: def with_policy(
      cls: type[module.Module], policy: Optional[jmp.Policy]
  ):
      def decorator(f):
          def wrapper(*args, **kwargs):
              with mixed_precision.push_policy(cls, policy):
                  return f(*args, **kwargs)
  
          return wrapper
  
      return decorator
  
imports: import importlib
  from typing import Optional
  from absl.testing import absltest
  from haiku._src import base
  from haiku._src import conv
  from haiku._src import mixed_precision
  from haiku._src import module
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import jmp
label_desc: Use image node_id 1 for calling a global function with example usage: with_policy(cls, policy) and returns: decorator, wrapper, f