Deploying to gh-pages from @ 7135fd7 🚀

_modules/blackjax/adaptation/mclmc_adaptation.html

@@ -391,7 +391,7 @@ <h1>Source code for blackjax.adaptation.mclmc_adaptation</h1><div class="highlig
 <span class="kn">from</span> <span class="nn">jax.flatten_util</span> <span class="kn">import</span> <span class="n">ravel_pytree</span>
 
 <span class="kn">from</span> <span class="nn">blackjax.diagnostics</span> <span class="kn">import</span> <span class="n">effective_sample_size</span>
-<span class="kn">from</span> <span class="nn">blackjax.util</span> <span class="kn">import</span> <span class="n">pytree_size</span><span class="p">,</span> <span class="n">streaming_average_update</span>
+<span class="kn">from</span> <span class="nn">blackjax.util</span> <span class="kn">import</span> <span class="n">incremental_value_update</span><span class="p">,</span> <span class="n">pytree_size</span>
 
 
 <div class="viewcode-block" id="MCLMCAdaptationState">
@@ -587,9 +587,9 @@ <h1>Source code for blackjax.adaptation.mclmc_adaptation</h1><div class="highlig
 
         <span class="n">x</span> <span class="o">=</span> <span class="n">ravel_pytree</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">position</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
         <span class="c1"># update the running average of x, x^2</span>
-        <span class="n">streaming_avg</span> <span class="o">=</span> <span class="n">streaming_average_update</span><span class="p">(</span>
-            <span class="n">current_value</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">x</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">square</span><span class="p">(</span><span class="n">x</span><span class="p">)]),</span>
-            <span class="n">previous_weight_and_average</span><span class="o">=</span><span class="n">streaming_avg</span><span class="p">,</span>
+        <span class="n">streaming_avg</span> <span class="o">=</span> <span class="n">incremental_value_update</span><span class="p">(</span>
+            <span class="n">expectation</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">x</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">square</span><span class="p">(</span><span class="n">x</span><span class="p">)]),</span>
+            <span class="n">incremental_val</span><span class="o">=</span><span class="n">streaming_avg</span><span class="p">,</span>
             <span class="n">weight</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">mask</span><span class="p">)</span> <span class="o">*</span> <span class="n">success</span> <span class="o">*</span> <span class="n">params</span><span class="o">.</span><span class="n">step_size</span><span class="p">,</span>
             <span class="n">zero_prevention</span><span class="o">=</span><span class="n">mask</span><span class="p">,</span>
         <span class="p">)</span>

examples/howto_custom_gradients.html

@@ -601,7 +601,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(1.2889438, dtype=float32, weak_type=True), logdensity=Array(-3.1444364, dtype=float32), logdensity_grad=Array(-2.3965275, dtype=float32, weak_type=True)))
+ HMCState(position=Array(1.2836552, dtype=float32, weak_type=True), logdensity=Array(-3.1318002, dtype=float32), logdensity_grad=Array(-2.3821688, dtype=float32, weak_type=True)))
 </pre></div>
 </div>
 </div>

examples/howto_metropolis_within_gibbs.html

@@ -607,8 +607,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 925 ms, sys: 38.1 ms, total: 963 ms
-Wall time: 954 ms
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 909 ms, sys: 47.8 ms, total: 957 ms
+Wall time: 946 ms
 </pre></div>
 </div>
 </div>
@@ -625,7 +625,7 @@ <h3>Sampling<a class="headerlink" href="#sampling" title="Link to this heading">
 </div>
 </div>
 <div class="cell_output docutils container">
-<img alt="../_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.png" src="../_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.png" />
+<img alt="../_images/4ea74d467c5e7fefec1f947e3b2d3435a1c39f9dd27aacd25cc898c44b621c29.png" src="../_images/4ea74d467c5e7fefec1f947e3b2d3435a1c39f9dd27aacd25cc898c44b621c29.png" />
 </div>
 </div>
 </section>
@@ -746,8 +746,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 700 ms, sys: 17.3 ms, total: 717 ms
-Wall time: 708 ms
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 723 ms, sys: 22.4 ms, total: 746 ms
+Wall time: 735 ms
 </pre></div>
 </div>
 </div>

examples/howto_sample_multiple_chains.html

@@ -527,8 +527,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 14.1 s, sys: 202 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.2 s, sys: 270 ms, total: 14.5 s
+Wall time: 14.5 s
 </pre></div>
 </div>
 </div>
@@ -632,8 +632,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 16.9 s, sys: 75.3 ms, total: 17 s
-Wall time: 5.56 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 18 s, sys: 59.4 ms, total: 18.1 s
+Wall time: 5.63 s
 </pre></div>
 </div>
 </div>

examples/howto_use_aesara.html

@@ -658,7 +658,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 0x7fde6584e7d0&gt;
+<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>&lt;ipykernel.iostream.OutStream at 0x7f36267ea7d0&gt;
 </pre></div>
 </div>
 </div>
@@ -760,7 +760,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/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png" src="../_images/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png" />
+<img alt="../_images/9636fcd531ffe29f41492f727e55f22fb0d0c83353777d312251e4923ec59532.png" src="../_images/9636fcd531ffe29f41492f727e55f22fb0d0c83353777d312251e4923ec59532.png" />
 </div>
 </div>
 </section>

examples/howto_use_numpyro.html

@@ -529,7 +529,7 @@ <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
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>\Average acceptance rate: 0.89
 There were 0.00% divergent transitions
 </pre></div>
 </div>
@@ -554,7 +554,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/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.png" src="../_images/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.png" />
+<img alt="../_images/1602b981d4682523edcb00799edc475a3156b5adb1cf42b7af0655100026fd9d.png" src="../_images/1602b981d4682523edcb00799edc475a3156b5adb1cf42b7af0655100026fd9d.png" />
 </div>
 </div>
 <div class="cell docutils container">
@@ -565,7 +565,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/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.png" src="../_images/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.png" />
+<img alt="../_images/f51b59090d070470bfdeb999de63513efb53c20a0163683e60c6bd15f70e7642.png" src="../_images/f51b59090d070470bfdeb999de63513efb53c20a0163683e60c6bd15f70e7642.png" />
 </div>
 </div>
 <div class="cell tag_hide-input docutils container">
@@ -584,14 +584,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.18
-Relative treatment effect for school 1: -0.13
-Relative treatment effect for school 2: 0.11
-Relative treatment effect for school 3: 0.04
-Relative treatment effect for school 4: 0.37
-Relative treatment effect for school 5: -0.28
-Relative treatment effect for school 6: 0.26
-Relative treatment effect for school 7: 0.26
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Relative treatment effect for school 0: 0.20
+Relative treatment effect for school 1: 0.44
+Relative treatment effect for school 2: -0.34
+Relative treatment effect for school 3: 0.14
+Relative treatment effect for school 4: -0.21
+Relative treatment effect for school 5: -0.13
+Relative treatment effect for school 6: -0.19
+Relative treatment effect for school 7: -0.25
 </pre></div>
 </div>
 </div>

examples/howto_use_oryx.html

@@ -549,8 +549,8 @@ <h1>Use with Oryx models<a class="headerlink" href="#use-with-oryx-models" 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.8 s, sys: 442 ms, total: 18.2 s
-Wall time: 18.1 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 18.2 s, sys: 379 ms, total: 18.6 s
+Wall time: 18.5 s
 </pre></div>
 </div>
 </div>
@@ -587,8 +587,8 @@ <h1>Use with Oryx models<a class="headerlink" href="#use-with-oryx-models" 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 16.3 s, sys: 163 ms, total: 16.5 s
-Wall time: 16.4 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 16.4 s, sys: 148 ms, total: 16.5 s
+Wall time: 16.5 s
 </pre></div>
 </div>
 </div>
@@ -629,7 +629,7 @@ <h1>Use with Oryx models<a class="headerlink" href="#use-with-oryx-models" title
 </div>
 </details>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average sample accuracy: 0.98080003
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average sample accuracy: 0.97933334
 BMA accuracy: 0.9866667
 </pre></div>
 </div>

examples/howto_use_pymc.html

@@ -514,7 +514,7 @@ <h1>Use with PyMC models<a class="headerlink" href="#use-with-pymc-models" title
 </div>
 </details>
 <div class="cell_output docutils container">
-<img alt="../_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png" src="../_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png" />
+<img alt="../_images/f91bbc45c5d240aebd6c757a6deb667dc93b88d42d64f3750feee22f179c55a3.png" src="../_images/f91bbc45c5d240aebd6c757a6deb667dc93b88d42d64f3750feee22f179c55a3.png" />
 </div>
 </div>
 </section>

examples/howto_use_tfp.html

@@ -511,7 +511,7 @@ <h1>Use with TFP models<a class="headerlink" href="#use-with-tfp-models" title="
 </div>
 </details>
 <div class="cell_output docutils container">
-<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average acceptance rate: 0.94
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average acceptance rate: 0.91
 </pre></div>
 </div>
 </div>
@@ -547,7 +547,7 @@ <h1>Use with TFP models<a class="headerlink" href="#use-with-tfp-models" title="
 </div>
 </details>
 <div class="cell_output docutils container">
-<img alt="../_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png" src="../_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png" />
+<img alt="../_images/10567ec44bdae7c0a3e89e32e7c7adcfa220f71eb3d35df75ea1fd7f10716ea6.png" src="../_images/10567ec44bdae7c0a3e89e32e7c7adcfa220f71eb3d35df75ea1fd7f10716ea6.png" />
 </div>
 </div>
 </section>

examples/quickstart.html

@@ -470,7 +470,7 @@ <h3>Set the Initial State<a class="headerlink" href="#set-the-initial-state" tit
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>HMCState(position={&#39;loc&#39;: 1.0, &#39;log_scale&#39;: 1.0}, logdensity=Array(-34471.617, dtype=float32), logdensity_grad={&#39;loc&#39;: Array(1216.7495, dtype=float32, weak_type=True), &#39;log_scale&#39;: Array(64105.36, dtype=float32, weak_type=True)})
+<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>HMCState(position={&#39;loc&#39;: 1.0, &#39;log_scale&#39;: 1.0}, logdensity=Array(-34937.934, dtype=float32), logdensity_grad={&#39;loc&#39;: Array(1158.9847, dtype=float32, weak_type=True), &#39;log_scale&#39;: Array(65037.996, dtype=float32, weak_type=True)})
 </pre></div>
 </div>
 </div>
@@ -518,8 +518,8 @@ <h3>Inference<a class="headerlink" href="#inference" 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 1.36 s, sys: 29.7 ms, total: 1.39 s
-Wall time: 1.38 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 1.37 s, sys: 54.1 ms, total: 1.43 s
+Wall time: 1.42 s
 </pre></div>
 </div>
 </div>
@@ -544,7 +544,7 @@ <h3>Inference<a class="headerlink" href="#inference" title="Link to this heading
 </div>
 </details>
 <div class="cell_output docutils container">
-<img alt="../_images/5375c0c0d96abe26f834f96ebf42a17a250ae7a0405b4e2df8055f3d9a7e9bd2.png" src="../_images/5375c0c0d96abe26f834f96ebf42a17a250ae7a0405b4e2df8055f3d9a7e9bd2.png" />
+<img alt="../_images/b081492c3e8202639cc929276d53a5b9fa16b687e2563c0cf10608f1010f4976.png" src="../_images/b081492c3e8202639cc929276d53a5b9fa16b687e2563c0cf10608f1010f4976.png" />
 </div>
 </div>
 </section>
@@ -571,7 +571,7 @@ <h2>NUTS<a class="headerlink" href="#nuts" title="Link to this heading">#</a></h
 </div>
 </div>
 <div class="cell_output docutils container">
-<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>HMCState(position={&#39;loc&#39;: 1.0, &#39;log_scale&#39;: 1.0}, logdensity=Array(-34471.617, dtype=float32), logdensity_grad={&#39;loc&#39;: Array(1216.7495, dtype=float32, weak_type=True), &#39;log_scale&#39;: Array(64105.36, dtype=float32, weak_type=True)})
+<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>HMCState(position={&#39;loc&#39;: 1.0, &#39;log_scale&#39;: 1.0}, logdensity=Array(-34937.934, dtype=float32), logdensity_grad={&#39;loc&#39;: Array(1158.9847, dtype=float32, weak_type=True), &#39;log_scale&#39;: Array(65037.996, dtype=float32, weak_type=True)})
 </pre></div>
 </div>
 </div>
@@ -588,8 +588,8 @@ <h2>NUTS<a class="headerlink" href="#nuts" title="Link to this heading">#</a></h
 </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 6.44 s, sys: 46.1 ms, total: 6.49 s
-Wall time: 6.47 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 6.22 s, sys: 72.6 ms, total: 6.29 s
+Wall time: 6.27 s
 </pre></div>
 </div>
 </div>
@@ -614,7 +614,7 @@ <h2>NUTS<a class="headerlink" href="#nuts" title="Link to this heading">#</a></h
 </div>
 </details>
 <div class="cell_output docutils container">
-<img alt="../_images/df5b789d507e8630887327b7d515b023c94aa10faccf135fbb3b63eeb3f6c81a.png" src="../_images/df5b789d507e8630887327b7d515b023c94aa10faccf135fbb3b63eeb3f6c81a.png" />
+<img alt="../_images/4fd93cad373e69c606a214e70efea94efdaf07910404ce95683e572310d7eb36.png" src="../_images/4fd93cad373e69c606a214e70efea94efdaf07910404ce95683e572310d7eb36.png" />
 </div>
 </div>
 <section id="use-stan-s-window-adaptation">
@@ -632,8 +632,8 @@ <h3>Use Stan’s Window Adaptation<a class="headerlink" href="#use-stan-s-window
 </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 1.36 s, sys: 81 ms, total: 1.44 s
-Wall time: 1.42 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 1.37 s, sys: 95.7 ms, total: 1.47 s
+Wall time: 1.45 s
 </pre></div>
 </div>
 </div>
@@ -652,8 +652,8 @@ <h3>Use Stan’s Window Adaptation<a class="headerlink" href="#use-stan-s-window
 </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 1 s, sys: 48.8 ms, total: 1.05 s
-Wall time: 1.03 s
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>CPU times: user 1.03 s, sys: 47.7 ms, total: 1.08 s
+Wall time: 1.06 s
 </pre></div>
 </div>
 </div>
@@ -678,7 +678,7 @@ <h3>Use Stan’s Window Adaptation<a class="headerlink" href="#use-stan-s-window
 </div>
 </details>
 <div class="cell_output docutils container">
-<img alt="../_images/5457d71d1aae97923d7f62a356f136f8584057e0cbc92755e0ea504730c0e4f9.png" src="../_images/5457d71d1aae97923d7f62a356f136f8584057e0cbc92755e0ea504730c0e4f9.png" />
+<img alt="../_images/7513574ee18cd1a2f22c2117bb87921b6e8f2d82f9536a8795d1fe1d828d1cdf.png" src="../_images/7513574ee18cd1a2f22c2117bb87921b6e8f2d82f9536a8795d1fe1d828d1cdf.png" />
 </div>
 </div>
 </section>