def METHOD_NAME(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
  )