code: def jax_binary_op(state, **kwargs):
      mkn = kwargs["mkn"]
      m = mkn[0]
      k = mkn[1]
      n = mkn[2]
      dtype = kwargs["dtype"]
      op = kwargs["op"]
      a = np.random.random([m, k]).astype(dtype)
      b = np.random.random([k, n]).astype(dtype)
      f = op
      f_jitted = jax.jit(f)
      f_jitted(a, b).block_until_ready()
      while state:
          f_jitted(a, b).block_until_ready()
      state.counters["items_per_second"] = Counter(
          state.iterations, Counter.kIsRate
      )
  
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 3 for calling a global function with example usage: jax_binary_op(state) without return types