code: def with_param_axis_error(f):
      @functools.wraps(f)
      def wrapper(*a, **k):
          old = layer_norm.ERROR_IF_PARAM_AXIS_NOT_EXPLICIT
          layer_norm.ERROR_IF_PARAM_AXIS_NOT_EXPLICIT = True
          try:
              return f(*a, **k)
          finally:
              layer_norm.ERROR_IF_PARAM_AXIS_NOT_EXPLICIT = old
  
      return wrapper
  
imports: import functools
  import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import initializers
  from haiku._src import layer_norm
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: with_param_axis_error(f) and returns: wrapper, f