Numpyro: Avoid compiling 2 times in HMC

[fehiepsi]

Currently, we do compiling 2 times: 1 at init, 1 at sample, so it would be better if we only compile 1 time.

One way is to modify tscan with additional parameters: skip to skip first warmup_steps; put warmup_state to a lax.cond to decide if it will be updated or not (depending on we are in warmup phase or sampling phase).

[neerajprad]

1 at init, 1 at sample, so it would be better if we only compile 1 time.

I think the compiled function is cached, so we shouldn't have to pay the cost the second time around, but lets verify this.

[neerajprad]

@fehiepsi - I think we can update this issue now to reflect the next steps on this task?

[fehiepsi]

It seems that we don't need to revise the api. It is likely an unexpected behaviour of lax.while_loop. More context about it can be found in https://github.com/google/jax/issues/587.

[neerajprad]

That's great! Should we close this issue then, or are there any things to follow up on?

[fehiepsi]

I think that it is better to let this open. In case it is not an issue of while_loop, I would prefer to add a small utility to do scan including warmup phase and avoid this issue.

[fehiepsi]

@neerajprad After working with the benchmark on covtype dataset, I think that we don't need to work around these issues of prims if we use nonprim version of lax.scan. Here are some benefits:

• Allows us force fast-compiling by feeding a test_state which makes hmc_next/nuts_next run fast
• Allows us separate compiling time and sampling time
• Allows us print/tqdm the sampling the progress

Here is a trade-off:

• We lost the benefit of using lax.scan/lax.fori_loop for the outer loop. But I expect that this is small.

Because we don't want the above trade-off create regressions on small model, let's address this issue again when more examples/benchmarks are added.

[neerajprad]

@fehiepsi - can we close this issue now?