def METHOD_NAME(state, **kwargs):
  shape = kwargs['shape']
  dtype = kwargs['dtype']
  op = kwargs['op']
  input0 = np.random.random(shape).astype(dtype)
  f = op
  f_jitted = jax.jit(f)
  f_jitted(input0).block_until_ready()
  while state:
    f_jitted(input0).block_until_ready()
  state.counters['items_per_second'] = Counter(
      input0.size * state.iterations, Counter.kIsRate
  )