Inner loop for benchmarking ar step.
(config, engine_generate, params, decode_state, iters, profile_name)
| 156 | |
| 157 | |
| 158 | def ar_benchmark_loop(config, engine_generate, params, decode_state, iters, profile_name): |
| 159 | """Inner loop for benchmarking ar step.""" |
| 160 | prof = profiler.Profiler(config) |
| 161 | prof.activate(optional_postfix=profile_name) |
| 162 | start = datetime.datetime.now() |
| 163 | rng = jax.random.PRNGKey(1234) |
| 164 | for _ in range(iters): |
| 165 | rng, rng_generate = jax.random.split(rng) |
| 166 | decode_state, _ = engine_generate(params, decode_state, rng_generate) |
| 167 | jax.block_until_ready(decode_state) |
| 168 | end = datetime.datetime.now() |
| 169 | prof.deactivate() |
| 170 | return (end - start).total_seconds(), decode_state |
| 171 | |
| 172 | |
| 173 | def ar_benchmark(config, engine_generate, params, decode_state, global_batch_size, cache_size, model_size, iters): |
no test coverage detected