code: def with_guardrails(f):
      """Runs the given test with JAX guardrails on."""
  
      @functools.wraps(f)
      def wrapper(*a, **k):
          old = config.check_jax_usage(True)
          try:
              return f(*a, **k)
          finally:
              config.check_jax_usage(old)
  
      return wrapper
  
file_path: dm-haiku/haiku/_src/test_utils.py
imports: from collections.abc import Generator, Sequence
  import functools
  import inspect
  import itertools
  import os
  import types
  from typing import Any, Callable, Optional, TypeVar
  from absl.testing import parameterized
  from haiku._src import config
  from haiku._src import transform
  import jax
label_desc: Use image node_id 8 for calling a global function with example usage: with_guardrails(f) and returns: wrapper, f