Deploying to gh-pages from @ 3f92393 🚀

blackjax-devs · May 8, 2024 · 826de87 · 826de87
1 parent 7633f8b
commit 826de87
Show file tree

Hide file tree

Showing 107 changed files with 212 additions and 56 deletions.
diff --git a/.doctrees/autoapi/blackjax/_version/index.doctree b/.doctrees/autoapi/blackjax/_version/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/base/index.doctree b/.doctrees/autoapi/blackjax/adaptation/base/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/chees_adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/chees_adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/mass_matrix/index.doctree b/.doctrees/autoapi/blackjax/adaptation/mass_matrix/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/mclmc_adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/mclmc_adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/meads_adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/meads_adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/pathfinder_adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/pathfinder_adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/step_size/index.doctree b/.doctrees/autoapi/blackjax/adaptation/step_size/index.doctree
diff --git a/.doctrees/autoapi/blackjax/adaptation/window_adaptation/index.doctree b/.doctrees/autoapi/blackjax/adaptation/window_adaptation/index.doctree
diff --git a/.doctrees/autoapi/blackjax/base/index.doctree b/.doctrees/autoapi/blackjax/base/index.doctree
diff --git a/.doctrees/autoapi/blackjax/diagnostics/index.doctree b/.doctrees/autoapi/blackjax/diagnostics/index.doctree
diff --git a/.doctrees/autoapi/blackjax/index.doctree b/.doctrees/autoapi/blackjax/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/barker/index.doctree b/.doctrees/autoapi/blackjax/mcmc/barker/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/diffusions/index.doctree b/.doctrees/autoapi/blackjax/mcmc/diffusions/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/dynamic_hmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/dynamic_hmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/elliptical_slice/index.doctree b/.doctrees/autoapi/blackjax/mcmc/elliptical_slice/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/ghmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/ghmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/hmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/hmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/integrators/index.doctree b/.doctrees/autoapi/blackjax/mcmc/integrators/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/mala/index.doctree b/.doctrees/autoapi/blackjax/mcmc/mala/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/marginal_latent_gaussian/index.doctree b/.doctrees/autoapi/blackjax/mcmc/marginal_latent_gaussian/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/mclmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/mclmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/metrics/index.doctree b/.doctrees/autoapi/blackjax/mcmc/metrics/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/nuts/index.doctree b/.doctrees/autoapi/blackjax/mcmc/nuts/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/periodic_orbital/index.doctree b/.doctrees/autoapi/blackjax/mcmc/periodic_orbital/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/proposal/index.doctree b/.doctrees/autoapi/blackjax/mcmc/proposal/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/random_walk/index.doctree b/.doctrees/autoapi/blackjax/mcmc/random_walk/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/rmhmc/index.doctree b/.doctrees/autoapi/blackjax/mcmc/rmhmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/termination/index.doctree b/.doctrees/autoapi/blackjax/mcmc/termination/index.doctree
diff --git a/.doctrees/autoapi/blackjax/mcmc/trajectory/index.doctree b/.doctrees/autoapi/blackjax/mcmc/trajectory/index.doctree
diff --git a/.doctrees/autoapi/blackjax/optimizers/dual_averaging/index.doctree b/.doctrees/autoapi/blackjax/optimizers/dual_averaging/index.doctree
diff --git a/.doctrees/autoapi/blackjax/optimizers/index.doctree b/.doctrees/autoapi/blackjax/optimizers/index.doctree
diff --git a/.doctrees/autoapi/blackjax/optimizers/lbfgs/index.doctree b/.doctrees/autoapi/blackjax/optimizers/lbfgs/index.doctree
diff --git a/.doctrees/autoapi/blackjax/progress_bar/index.doctree b/.doctrees/autoapi/blackjax/progress_bar/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/csgld/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/csgld/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/diffusions/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/diffusions/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/gradients/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/gradients/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/sghmc/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/sghmc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/sgld/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/sgld/index.doctree
diff --git a/.doctrees/autoapi/blackjax/sgmcmc/sgnht/index.doctree b/.doctrees/autoapi/blackjax/sgmcmc/sgnht/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/adaptive_tempered/index.doctree b/.doctrees/autoapi/blackjax/smc/adaptive_tempered/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/base/index.doctree b/.doctrees/autoapi/blackjax/smc/base/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/ess/index.doctree b/.doctrees/autoapi/blackjax/smc/ess/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/index.doctree b/.doctrees/autoapi/blackjax/smc/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/inner_kernel_tuning/index.doctree b/.doctrees/autoapi/blackjax/smc/inner_kernel_tuning/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/resampling/index.doctree b/.doctrees/autoapi/blackjax/smc/resampling/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/solver/index.doctree b/.doctrees/autoapi/blackjax/smc/solver/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/tempered/index.doctree b/.doctrees/autoapi/blackjax/smc/tempered/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/tuning/from_kernel_info/index.doctree b/.doctrees/autoapi/blackjax/smc/tuning/from_kernel_info/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/tuning/from_particles/index.doctree b/.doctrees/autoapi/blackjax/smc/tuning/from_particles/index.doctree
diff --git a/.doctrees/autoapi/blackjax/smc/tuning/index.doctree b/.doctrees/autoapi/blackjax/smc/tuning/index.doctree
diff --git a/.doctrees/autoapi/blackjax/types/index.doctree b/.doctrees/autoapi/blackjax/types/index.doctree
diff --git a/.doctrees/autoapi/blackjax/util/index.doctree b/.doctrees/autoapi/blackjax/util/index.doctree
diff --git a/.doctrees/autoapi/blackjax/vi/index.doctree b/.doctrees/autoapi/blackjax/vi/index.doctree
diff --git a/.doctrees/autoapi/blackjax/vi/meanfield_vi/index.doctree b/.doctrees/autoapi/blackjax/vi/meanfield_vi/index.doctree
diff --git a/.doctrees/autoapi/blackjax/vi/pathfinder/index.doctree b/.doctrees/autoapi/blackjax/vi/pathfinder/index.doctree
diff --git a/.doctrees/autoapi/blackjax/vi/schrodinger_follmer/index.doctree b/.doctrees/autoapi/blackjax/vi/schrodinger_follmer/index.doctree
diff --git a/.doctrees/autoapi/blackjax/vi/svgd/index.doctree b/.doctrees/autoapi/blackjax/vi/svgd/index.doctree
diff --git a/.doctrees/bib.doctree b/.doctrees/bib.doctree
diff --git a/.doctrees/environment.pickle b/.doctrees/environment.pickle
diff --git a/.doctrees/examples/howto_custom_gradients.doctree b/.doctrees/examples/howto_custom_gradients.doctree
diff --git a/.doctrees/examples/howto_metropolis_within_gibbs.doctree b/.doctrees/examples/howto_metropolis_within_gibbs.doctree
diff --git a/.doctrees/examples/howto_other_frameworks.doctree b/.doctrees/examples/howto_other_frameworks.doctree
diff --git a/.doctrees/examples/howto_reproduce_the_blackjax_image.doctree b/.doctrees/examples/howto_reproduce_the_blackjax_image.doctree
diff --git a/.doctrees/examples/howto_sample_multiple_chains.doctree b/.doctrees/examples/howto_sample_multiple_chains.doctree
diff --git a/.doctrees/examples/howto_use_aesara.doctree b/.doctrees/examples/howto_use_aesara.doctree
diff --git a/.doctrees/examples/howto_use_numpyro.doctree b/.doctrees/examples/howto_use_numpyro.doctree
diff --git a/.doctrees/examples/howto_use_oryx.doctree b/.doctrees/examples/howto_use_oryx.doctree
diff --git a/.doctrees/examples/howto_use_pymc.doctree b/.doctrees/examples/howto_use_pymc.doctree
diff --git a/.doctrees/examples/howto_use_tfp.doctree b/.doctrees/examples/howto_use_tfp.doctree
diff --git a/.doctrees/examples/quickstart.doctree b/.doctrees/examples/quickstart.doctree
diff --git a/.doctrees/index.doctree b/.doctrees/index.doctree
diff --git a/_images/08f150c1d1f87ea47eb3741da69fb39b18cbd55912b2048e5297493936fd53b2.png b/_images/08f150c1d1f87ea47eb3741da69fb39b18cbd55912b2048e5297493936fd53b2.png
diff --git a/_images/2e01b15794a2a0a0c9e03b26c69423105e4f4840ce200b9225abbc2cfe3540fb.png b/_images/2e01b15794a2a0a0c9e03b26c69423105e4f4840ce200b9225abbc2cfe3540fb.png
diff --git a/_images/30e7934c691429575f48d98ef7b38071b9884e6152859891c3035a6cfa769878.png b/_images/30e7934c691429575f48d98ef7b38071b9884e6152859891c3035a6cfa769878.png
diff --git a/_images/4244d9691b22d64aab50867efb7ad8a866ecf061c2529f56af798f429961f2f5.png b/_images/4244d9691b22d64aab50867efb7ad8a866ecf061c2529f56af798f429961f2f5.png
diff --git a/_images/44957ea8cad98d0d4cc3430cd14d927ebd4fc00670f4be34236be661c70a2921.png b/_images/44957ea8cad98d0d4cc3430cd14d927ebd4fc00670f4be34236be661c70a2921.png
diff --git a/_images/4a6ec46b1cf51a550920214e7085fffccebffa359d663ff67223ec290a9ae507.png b/_images/4a6ec46b1cf51a550920214e7085fffccebffa359d663ff67223ec290a9ae507.png
diff --git a/_images/4c898299e157f1af4339ed14bd05cfa55e78d23518c022af288183ae88ef3545.png b/_images/4c898299e157f1af4339ed14bd05cfa55e78d23518c022af288183ae88ef3545.png
diff --git a/_images/511d92e85bc2aab534fd90fb7fe045b0ab25d131a45052e92f6b2c9a97925f20.png b/_images/511d92e85bc2aab534fd90fb7fe045b0ab25d131a45052e92f6b2c9a97925f20.png
diff --git a/_images/5e9dacb1cc8432322d836df7671d6f1d06adbfd64dab4d9e339ef9c9cc642a65.png b/_images/5e9dacb1cc8432322d836df7671d6f1d06adbfd64dab4d9e339ef9c9cc642a65.png
diff --git a/_images/69ed75d9b6fd6050d7ae4c51c42fd0711de160e2ba2d0fe3f508fba1bed3d104.png b/_images/69ed75d9b6fd6050d7ae4c51c42fd0711de160e2ba2d0fe3f508fba1bed3d104.png
diff --git a/_images/80afb178ff06ce447f9d6999e3a9d1e66b86f590a85cbcc288811a872af9e9ac.png b/_images/80afb178ff06ce447f9d6999e3a9d1e66b86f590a85cbcc288811a872af9e9ac.png
diff --git a/_images/87f252410236e8e00892ad49f60db145376a4f29182f7b2dff2b2cbe5011c9c7.png b/_images/87f252410236e8e00892ad49f60db145376a4f29182f7b2dff2b2cbe5011c9c7.png
diff --git a/_images/8f62409cb41460baa7e2ddb51470f9979ec2d121ac59e81fcb8fe2b232976686.png b/_images/8f62409cb41460baa7e2ddb51470f9979ec2d121ac59e81fcb8fe2b232976686.png
diff --git a/_images/9b409e3e741c913e0ab9f0ccebd7198c8f934909300e740c767aa3dede642679.png b/_images/9b409e3e741c913e0ab9f0ccebd7198c8f934909300e740c767aa3dede642679.png
diff --git a/_images/aa8f91c89c3588232511e4a3bfb7e0f138ecc96be0bc378d707d58f3bca649b6.png b/_images/aa8f91c89c3588232511e4a3bfb7e0f138ecc96be0bc378d707d58f3bca649b6.png
diff --git a/_images/aac3c058972d7d5665f82fcd887f08146aefa40920da82a70d8e5cf03821a268.png b/_images/aac3c058972d7d5665f82fcd887f08146aefa40920da82a70d8e5cf03821a268.png
diff --git a/_images/baffa78b233e62eaae0b452fedc34dd869b6a850dcf567c3bb47d0bcc867b97c.png b/_images/baffa78b233e62eaae0b452fedc34dd869b6a850dcf567c3bb47d0bcc867b97c.png
diff --git a/_images/c1b5ae4d60046bbc7df177257e50d8886674d8e044040b7097299f2b53263d7c.png b/_images/c1b5ae4d60046bbc7df177257e50d8886674d8e044040b7097299f2b53263d7c.png
diff --git a/_modules/blackjax/_version.html b/_modules/blackjax/_version.html
@@ -392,8 +392,8 @@ <h1>Source code for blackjax._version</h1><div class="highlight"><pre>
 <span class="n">version_tuple</span><span class="p">:</span> <span class="n">VERSION_TUPLE</span></div>
 
 
-<span class="n">__version__</span> <span class="o">=</span> <span class="n">version</span> <span class="o">=</span> <span class="s1">&#39;0.1.dev1+g1bc6f93&#39;</span>
-<span class="n">__version_tuple__</span> <span class="o">=</span> <span class="n">version_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="s1">&#39;dev1&#39;</span><span class="p">,</span> <span class="s1">&#39;g1bc6f93&#39;</span><span class="p">)</span>
+<span class="n">__version__</span> <span class="o">=</span> <span class="n">version</span> <span class="o">=</span> <span class="s1">&#39;0.1.dev1+g3f92393&#39;</span>
+<span class="n">__version_tuple__</span> <span class="o">=</span> <span class="n">version_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="s1">&#39;dev1&#39;</span><span class="p">,</span> <span class="s1">&#39;g3f92393&#39;</span><span class="p">)</span>
 </pre></div>
 
                 </article>

diff --git a/_modules/blackjax/mcmc/proposal.html b/_modules/blackjax/mcmc/proposal.html
@@ -550,7 +550,7 @@ <h1>Source code for blackjax.mcmc.proposal</h1><div class="highlight"><pre>
 <span class="sd">    biases the transition away from the trajectory&#39;s initial state.</span>
 
 <span class="sd">    &quot;&quot;&quot;</span>
-    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">new_proposal</span><span class="o">.</span><span class="n">weight</span> <span class="o">-</span> <span class="n">proposal</span><span class="o">.</span><span class="n">weight</span><span class="p">),</span> <span class="n">a_max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
+    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">new_proposal</span><span class="o">.</span><span class="n">weight</span> <span class="o">-</span> <span class="n">proposal</span><span class="o">.</span><span class="n">weight</span><span class="p">),</span> <span class="nb">max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
     <span class="n">do_accept</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">bernoulli</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">p_accept</span><span class="p">)</span>
     <span class="n">new_weight</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">logaddexp</span><span class="p">(</span><span class="n">proposal</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">new_proposal</span><span class="o">.</span><span class="n">weight</span><span class="p">)</span>
     <span class="n">new_sum_log_p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">logaddexp</span><span class="p">(</span>
@@ -627,7 +627,7 @@ <h1>Source code for blackjax.mcmc.proposal</h1><div class="highlight"><pre>
 <span class="sd">    then the new proposal is accepted with probability 1.</span>
 
 <span class="sd">    &quot;&quot;&quot;</span>
-    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">log_p_accept</span><span class="p">),</span> <span class="n">a_max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
+    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">log_p_accept</span><span class="p">),</span> <span class="nb">max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
     <span class="n">do_accept</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">bernoulli</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">p_accept</span><span class="p">)</span>
     <span class="n">info</span> <span class="o">=</span> <span class="n">do_accept</span><span class="p">,</span> <span class="n">p_accept</span><span class="p">,</span> <span class="kc">None</span>
     <span class="k">return</span> <span class="p">(</span>
@@ -659,7 +659,7 @@ <h1>Source code for blackjax.mcmc.proposal</h1><div class="highlight"><pre>
 <span class="sd">    to the accept/reject step of a current state and new proposal.</span>
 
 <span class="sd">    &quot;&quot;&quot;</span>
-    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">delta_energy</span><span class="p">),</span> <span class="n">a_max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
+    <span class="n">p_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">delta_energy</span><span class="p">),</span> <span class="nb">max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
     <span class="n">do_accept</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="nb">slice</span><span class="p">))</span> <span class="o">&lt;=</span> <span class="n">delta_energy</span>
     <span class="n">slice_next</span> <span class="o">=</span> <span class="nb">slice</span> <span class="o">*</span> <span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="n">delta_energy</span><span class="p">)</span> <span class="o">*</span> <span class="n">do_accept</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">do_accept</span><span class="p">))</span>
     <span class="n">info</span> <span class="o">=</span> <span class="n">do_accept</span><span class="p">,</span> <span class="n">p_accept</span><span class="p">,</span> <span class="n">slice_next</span>

diff --git a/examples/howto_custom_gradients.html b/examples/howto_custom_gradients.html
@@ -599,7 +599,7 @@ <h3>Using the function with Blackjax<a class="headerlink" href="#using-the-funct
 </div>
 <div class="cell_output docutils container">
 <div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>(HMCState(position=1.0, logdensity=Array(-2.5600991, dtype=float32), logdensity_grad=Array(-1.6666666, dtype=float32, weak_type=True)),
- HMCState(position=Array(0.8047719, dtype=float32, weak_type=True), logdensity=Array(-2.277532, dtype=float32), logdensity_grad=Array(-1.2365438, dtype=float32, weak_type=True)))
+ HMCState(position=Array(0.97067314, dtype=float32, weak_type=True), logdensity=Array(-2.512219, dtype=float32), logdensity_grad=Array(-1.5988109, dtype=float32, weak_type=True)))
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_metropolis_within_gibbs.html b/examples/howto_metropolis_within_gibbs.html
@@ -604,8 +604,8 @@ <h3>Sampling<a class="headerlink" href="#sampling" title="Link to this heading">
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 951 ms, sys: 61.8 ms, total: 1.01 s
-Wall time: 1.01 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 919 ms, sys: 67.5 ms, total: 986 ms
+Wall time: 980 ms
 </pre></div>
 </div>
 </div>
@@ -622,7 +622,7 @@ <h3>Sampling<a class="headerlink" href="#sampling" title="Link to this heading">
 </div>
 </div>
 <div class="cell_output docutils container">
-<img alt="../_images/80afb178ff06ce447f9d6999e3a9d1e66b86f590a85cbcc288811a872af9e9ac.png" src="../_images/80afb178ff06ce447f9d6999e3a9d1e66b86f590a85cbcc288811a872af9e9ac.png" />
+<img alt="../_images/aa8f91c89c3588232511e4a3bfb7e0f138ecc96be0bc378d707d58f3bca649b6.png" src="../_images/aa8f91c89c3588232511e4a3bfb7e0f138ecc96be0bc378d707d58f3bca649b6.png" />
 </div>
 </div>
 </section>
@@ -743,8 +743,8 @@ <h3>Sampling<a class="headerlink" href="#id2" title="Link to this heading">#</a>
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 696 ms, sys: 16.8 ms, total: 713 ms
-Wall time: 703 ms
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 705 ms, sys: 36.5 ms, total: 742 ms
+Wall time: 733 ms
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_sample_multiple_chains.html b/examples/howto_sample_multiple_chains.html
@@ -524,8 +524,8 @@ <h3>Using <code class="docutils literal notranslate"><span class="pre">jax.vmap<
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 13.8 s, sys: 508 ms, total: 14.3 s
-Wall time: 14.3 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 14.9 s, sys: 201 ms, total: 15.1 s
+Wall time: 15.1 s
 </pre></div>
 </div>
 </div>
@@ -629,8 +629,8 @@ <h3>Back to our example<a class="headerlink" href="#back-to-our-example" title="
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 17.7 s, sys: 60.1 ms, total: 17.8 s
-Wall time: 5.69 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 17.9 s, sys: 76.3 ms, total: 18 s
+Wall time: 5.62 s
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_use_aesara.html b/examples/howto_use_aesara.html
@@ -655,7 +655,7 @@ <h1>Use with Aesara models<a class="headerlink" href="#use-with-aesara-models" t
                |Elemwise{sigmoid,no_inplace} [id FD]
 </pre></div>
 </div>
-<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>&lt;ipykernel.iostream.OutStream at 0x7fa9e50c05b0&gt;
+<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>&lt;ipykernel.iostream.OutStream at 0x7fe2f43185b0&gt;
 </pre></div>
 </div>
 </div>
@@ -757,7 +757,7 @@ <h1>Use with Aesara models<a class="headerlink" href="#use-with-aesara-models" t
 </div>
 </div>
 <div class="cell_output docutils container">
-<img alt="../_images/4a6ec46b1cf51a550920214e7085fffccebffa359d663ff67223ec290a9ae507.png" src="../_images/4a6ec46b1cf51a550920214e7085fffccebffa359d663ff67223ec290a9ae507.png" />
+<img alt="../_images/baffa78b233e62eaae0b452fedc34dd869b6a850dcf567c3bb47d0bcc867b97c.png" src="../_images/baffa78b233e62eaae0b452fedc34dd869b6a850dcf567c3bb47d0bcc867b97c.png" />
 </div>
 </div>
 </section>

diff --git a/examples/howto_use_numpyro.html b/examples/howto_use_numpyro.html
@@ -526,8 +526,8 @@ <h1>Use with Numpyro models<a class="headerlink" href="#use-with-numpyro-models"
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>\Average acceptance rate: 0.93
-There were 0.00% divergent transitions
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>\Average acceptance rate: 0.86
+There were 0.10% divergent transitions
 </pre></div>
 </div>
 </div>
@@ -551,7 +551,7 @@ <h1>Use with Numpyro models<a class="headerlink" href="#use-with-numpyro-models"
 </div>
 </div>
 <div class="cell_output docutils container">
-<img alt="../_images/5e9dacb1cc8432322d836df7671d6f1d06adbfd64dab4d9e339ef9c9cc642a65.png" src="../_images/5e9dacb1cc8432322d836df7671d6f1d06adbfd64dab4d9e339ef9c9cc642a65.png" />
+<img alt="../_images/c1b5ae4d60046bbc7df177257e50d8886674d8e044040b7097299f2b53263d7c.png" src="../_images/c1b5ae4d60046bbc7df177257e50d8886674d8e044040b7097299f2b53263d7c.png" />
 </div>
 </div>
 <div class="cell docutils container">
@@ -562,7 +562,7 @@ <h1>Use with Numpyro models<a class="headerlink" href="#use-with-numpyro-models"
 </div>
 </div>
 <div class="cell_output docutils container">
-<img alt="../_images/69ed75d9b6fd6050d7ae4c51c42fd0711de160e2ba2d0fe3f508fba1bed3d104.png" src="../_images/69ed75d9b6fd6050d7ae4c51c42fd0711de160e2ba2d0fe3f508fba1bed3d104.png" />
+<img alt="../_images/44957ea8cad98d0d4cc3430cd14d927ebd4fc00670f4be34236be661c70a2921.png" src="../_images/44957ea8cad98d0d4cc3430cd14d927ebd4fc00670f4be34236be661c70a2921.png" />
 </div>
 </div>
 <div class="cell tag_hide-input docutils container">
@@ -581,14 +581,14 @@ <h1>Use with Numpyro models<a class="headerlink" href="#use-with-numpyro-models"
 </div>
 </details>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Relative treatment effect for school 0: 0.31
-Relative treatment effect for school 1: 0.56
-Relative treatment effect for school 2: 0.07
-Relative treatment effect for school 3: -0.11
-Relative treatment effect for school 4: 0.60
-Relative treatment effect for school 5: 0.47
-Relative treatment effect for school 6: 0.23
-Relative treatment effect for school 7: -0.16
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Relative treatment effect for school 0: 0.32
+Relative treatment effect for school 1: 0.41
+Relative treatment effect for school 2: -0.27
+Relative treatment effect for school 3: -0.27
+Relative treatment effect for school 4: 0.09
+Relative treatment effect for school 5: -0.01
+Relative treatment effect for school 6: 0.33
+Relative treatment effect for school 7: 0.12
 </pre></div>
 </div>
 </div>