code: def math_benchmark(*args):
      def decorator(func):
          for test_case in args[0]:
  
              @benchmark.register(
                  name=f"{func.__name__}_{test_case['name']}"
              )
              @functools.wraps(func)
              def wrapper(state, test_case=test_case):
                  return func(state, **test_case)
  
          return wrapper
  
      return decorator
  
file_path: jax/benchmarks/math_benchmark.py
imports: import functools
  import google_benchmark
  import jax
  import jax.numpy
  import numpy
  import os
  import sys
  from google_benchmark import Counter
  from absl import app
  from absl import flags
label_desc: Use image node_id 1 for calling a global function with example usage: math_benchmark() and returns: decorator, wrapper, func