Update use_with_numpyro.ipynb

[junpenglao]

Update notebook to get the same diagnostic information from both numpyro and blackjax NUTS, and also setting some parameters to be the same.

This also help benchmarking as now both sampler ran in the same parameter setting.

[junpenglao]

To use this notebook as simple benchmark, i first set the flag in the notebook RUN_BENCHMARK = True, then from terminal:

(blackjax) junpenglao-macbookpro% ipython examples/use_with_numpyro.ipynb

Benchmark with 1000 warmup steps and 5000000 sampling steps.

Blackjax:
-> Running warmup.
CPU times: user 3.06 s, sys: 48.9 ms, total: 3.11 s
Wall time: 3.12 s
-> Running sampling.
CPU times: user 29.8 s, sys: 673 ms, total: 30.4 s
Wall time: 30.1 s

Acceptance rate: 0.87
0.04% divergent transitions

Numpyro:
-> Running warmup+sampling.
CPU times: user 30.9 s, sys: 500 ms, total: 31.4 s
Wall time: 30.9 s

Acceptance rate: 0.90
0.00% divergent transitions

Blackjax average 12.08 leapfrog per iteration.
Numpyro average 10.86 leapfrog per iteration.

[codecov]

Codecov Report

Merging #165 (ceccdf8) into main (70e6b19) will not change coverage.
The diff coverage is n/a.

❗ Current head ceccdf8 differs from pull request most recent head 58e46eb. Consider uploading reports for the commit 58e46eb to get more accurate results

@@           Coverage Diff           @@
##             main     #165   +/-   ##
=======================================
  Coverage   98.56%   98.56%           
=======================================
  Files          25       25           
  Lines         973      973           
=======================================
  Hits          959      959           
  Misses         14       14           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 70e6b19...58e46eb. Read the comment docs.

[rlouf]

Clearly the window adaptation is not optimal, cf #44. But we can already see blackjax is faster / gradient evaluation :)