Deploying to gh-pages from @ 148c028 🚀

blackjax-devs · Aug 9, 2024 · 475bd57 · 475bd57
1 parent 0914121
commit 475bd57
Show file tree

Hide file tree

Showing 40 changed files with 48 additions and 55 deletions.
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_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/_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png b/_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png
diff --git a/_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.png b/_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.png
diff --git a/_images/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.png b/_images/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.png
diff --git a/_images/3f151773ff7adbbca9bc0ab72466bb701ae174029b2e01f3707df3529df580d4.png b/_images/3f151773ff7adbbca9bc0ab72466bb701ae174029b2e01f3707df3529df580d4.png
diff --git a/_images/5375c0c0d96abe26f834f96ebf42a17a250ae7a0405b4e2df8055f3d9a7e9bd2.png b/_images/5375c0c0d96abe26f834f96ebf42a17a250ae7a0405b4e2df8055f3d9a7e9bd2.png
diff --git a/_images/5457d71d1aae97923d7f62a356f136f8584057e0cbc92755e0ea504730c0e4f9.png b/_images/5457d71d1aae97923d7f62a356f136f8584057e0cbc92755e0ea504730c0e4f9.png
diff --git a/_images/57d7f53d35f017a4df17406fedc1fdca63d01600ec2de38d0be5081474b4f09d.png b/_images/57d7f53d35f017a4df17406fedc1fdca63d01600ec2de38d0be5081474b4f09d.png
diff --git a/_images/607983e7d9ae823d53d48df672ad227d7606464223954292df409a88bf3f1584.png b/_images/607983e7d9ae823d53d48df672ad227d7606464223954292df409a88bf3f1584.png
diff --git a/_images/63ce52ef0a43e13ae595467d6b83735f618d59def295f91b1510661180de9484.png b/_images/63ce52ef0a43e13ae595467d6b83735f618d59def295f91b1510661180de9484.png
diff --git a/_images/6641e0a1f461ab4f209f16e14765a2c90c37b0f88e84630ae12736f57730d3ef.png b/_images/6641e0a1f461ab4f209f16e14765a2c90c37b0f88e84630ae12736f57730d3ef.png
diff --git a/_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png b/_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png
diff --git a/_images/9efd18b18e29afd1a43e8bee4324236ed9f025c2cf5a44f53b1e7b5f0f66d722.png b/_images/9efd18b18e29afd1a43e8bee4324236ed9f025c2cf5a44f53b1e7b5f0f66d722.png
diff --git a/_images/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.png b/_images/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.png
diff --git a/_images/a336b164822b7973ab73802611ae9b0d156b7801425cf7c2e366c3a09e1bc57a.png b/_images/a336b164822b7973ab73802611ae9b0d156b7801425cf7c2e366c3a09e1bc57a.png
diff --git a/_images/adc1e7538b93ddc1114cc816750888353642cbd9f3d50883a92c3935f9d93f35.png b/_images/adc1e7538b93ddc1114cc816750888353642cbd9f3d50883a92c3935f9d93f35.png
diff --git a/_images/df5b789d507e8630887327b7d515b023c94aa10faccf135fbb3b63eeb3f6c81a.png b/_images/df5b789d507e8630887327b7d515b023c94aa10faccf135fbb3b63eeb3f6c81a.png
diff --git a/_images/e2dc07d29cae908bec94d4b847d8b7e400b38641709a08ce3cf3dd89dee7143d.png b/_images/e2dc07d29cae908bec94d4b847d8b7e400b38641709a08ce3cf3dd89dee7143d.png
diff --git a/_images/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png b/_images/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png
diff --git a/_modules/blackjax/adaptation/window_adaptation.html b/_modules/blackjax/adaptation/window_adaptation.html
@@ -399,7 +399,7 @@ <h1>Source code for blackjax.adaptation.window_adaptation</h1><div class="highli
     <span class="n">dual_averaging_adaptation</span><span class="p">,</span>
 <span class="p">)</span>
 <span class="kn">from</span> <span class="nn">blackjax.base</span> <span class="kn">import</span> <span class="n">AdaptationAlgorithm</span>
-<span class="kn">from</span> <span class="nn">blackjax.progress_bar</span> <span class="kn">import</span> <span class="n">progress_bar_scan</span>
+<span class="kn">from</span> <span class="nn">blackjax.progress_bar</span> <span class="kn">import</span> <span class="n">gen_scan_fn</span>
 <span class="kn">from</span> <span class="nn">blackjax.types</span> <span class="kn">import</span> <span class="n">Array</span><span class="p">,</span> <span class="n">ArrayLikeTree</span><span class="p">,</span> <span class="n">PRNGKey</span>
 <span class="kn">from</span> <span class="nn">blackjax.util</span> <span class="kn">import</span> <span class="n">pytree_size</span>
 
@@ -724,23 +724,16 @@ <h1>Source code for blackjax.adaptation.window_adaptation</h1><div class="highli
 
         <span class="k">if</span> <span class="n">progress_bar</span><span class="p">:</span>
             <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Running window adaptation&quot;</span><span class="p">)</span>
-            <span class="n">one_step_</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">jit</span><span class="p">(</span><span class="n">progress_bar_scan</span><span class="p">(</span><span class="n">num_steps</span><span class="p">)(</span><span class="n">one_step</span><span class="p">))</span>
-            <span class="n">start_state</span> <span class="o">=</span> <span class="p">((</span><span class="n">init_state</span><span class="p">,</span> <span class="n">init_adaptation_state</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
-        <span class="k">else</span><span class="p">:</span>
-            <span class="n">one_step_</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">jit</span><span class="p">(</span><span class="n">one_step</span><span class="p">)</span>
-            <span class="n">start_state</span> <span class="o">=</span> <span class="p">(</span><span class="n">init_state</span><span class="p">,</span> <span class="n">init_adaptation_state</span><span class="p">)</span>
-
+        <span class="n">scan_fn</span> <span class="o">=</span> <span class="n">gen_scan_fn</span><span class="p">(</span><span class="n">num_steps</span><span class="p">,</span> <span class="n">progress_bar</span><span class="o">=</span><span class="n">progress_bar</span><span class="p">)</span>
+        <span class="n">start_state</span> <span class="o">=</span> <span class="p">(</span><span class="n">init_state</span><span class="p">,</span> <span class="n">init_adaptation_state</span><span class="p">)</span>
         <span class="n">keys</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">split</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">num_steps</span><span class="p">)</span>
         <span class="n">schedule</span> <span class="o">=</span> <span class="n">build_schedule</span><span class="p">(</span><span class="n">num_steps</span><span class="p">)</span>
-        <span class="n">last_state</span><span class="p">,</span> <span class="n">info</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">lax</span><span class="o">.</span><span class="n">scan</span><span class="p">(</span>
-            <span class="n">one_step_</span><span class="p">,</span>
+        <span class="n">last_state</span><span class="p">,</span> <span class="n">info</span> <span class="o">=</span> <span class="n">scan_fn</span><span class="p">(</span>
+            <span class="n">one_step</span><span class="p">,</span>
             <span class="n">start_state</span><span class="p">,</span>
             <span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">num_steps</span><span class="p">),</span> <span class="n">keys</span><span class="p">,</span> <span class="n">schedule</span><span class="p">),</span>
         <span class="p">)</span>
 
-        <span class="k">if</span> <span class="n">progress_bar</span><span class="p">:</span>
-            <span class="n">last_state</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">last_state</span>
-
         <span class="n">last_chain_state</span><span class="p">,</span> <span class="n">last_warmup_state</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span> <span class="o">=</span> <span class="n">last_state</span>
 
         <span class="n">step_size</span><span class="p">,</span> <span class="n">inverse_mass_matrix</span> <span class="o">=</span> <span class="n">adapt_final</span><span class="p">(</span><span class="n">last_warmup_state</span><span class="p">)</span>

diff --git a/examples/howto_custom_gradients.html b/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(0.65200007, dtype=float32, weak_type=True), logdensity=Array(-2.112022, dtype=float32), logdensity_grad=Array(-0.9354014, 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)))
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_metropolis_within_gibbs.html b/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 895 ms, sys: 52.3 ms, total: 947 ms
-Wall time: 937 ms
+<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
 </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/6641e0a1f461ab4f209f16e14765a2c90c37b0f88e84630ae12736f57730d3ef.png" src="../_images/6641e0a1f461ab4f209f16e14765a2c90c37b0f88e84630ae12736f57730d3ef.png" />
+<img alt="../_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.png" src="../_images/1d7c9a6b8b932c078865f843202a575a9f7d4317ee94589b85c55c228f734f4b.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 702 ms, sys: 12.5 ms, total: 714 ms
-Wall time: 705 ms
+<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
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_sample_multiple_chains.html b/examples/howto_sample_multiple_chains.html
@@ -527,7 +527,7 @@ <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: 238 ms, total: 14.4 s
+<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
 </pre></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 17.2 s, sys: 51.6 ms, total: 17.2 s
-Wall time: 5.5 s
+<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
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_use_aesara.html b/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 0x7f47fd31a830&gt;
+<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>&lt;ipykernel.iostream.OutStream at 0x7fde6584e7d0&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/9efd18b18e29afd1a43e8bee4324236ed9f025c2cf5a44f53b1e7b5f0f66d722.png" src="../_images/9efd18b18e29afd1a43e8bee4324236ed9f025c2cf5a44f53b1e7b5f0f66d722.png" />
+<img alt="../_images/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png" src="../_images/f60107edb02c8b3fb7450218f9d87727932635df0224a3f5a00d8760b0749ff3.png" />
 </div>
 </div>
 </section>

diff --git a/examples/howto_use_numpyro.html b/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.85
+<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
 </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/57d7f53d35f017a4df17406fedc1fdca63d01600ec2de38d0be5081474b4f09d.png" src="../_images/57d7f53d35f017a4df17406fedc1fdca63d01600ec2de38d0be5081474b4f09d.png" />
+<img alt="../_images/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.png" src="../_images/270f58a2bd1b5e3225c3646db807868113a85637c7d7b976d676c6e83148ae88.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/63ce52ef0a43e13ae595467d6b83735f618d59def295f91b1510661180de9484.png" src="../_images/63ce52ef0a43e13ae595467d6b83735f618d59def295f91b1510661180de9484.png" />
+<img alt="../_images/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.png" src="../_images/a06c93a909f0a41e9eba9f25f60d9d3fd143f5c184d786b09bc6b678c1a81486.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.51
-Relative treatment effect for school 1: -0.17
-Relative treatment effect for school 2: -0.02
-Relative treatment effect for school 3: -0.26
-Relative treatment effect for school 4: -0.30
-Relative treatment effect for school 5: 0.66
-Relative treatment effect for school 6: 0.40
-Relative treatment effect for school 7: -0.19
+<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
 </pre></div>
 </div>
 </div>

diff --git a/examples/howto_use_oryx.html b/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 18.2 s, sys: 371 ms, total: 18.6 s
-Wall time: 18.5 s
+<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
 </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: 147 ms, total: 16.4 s
-Wall time: 16.3 s
+<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
 </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.981
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average sample accuracy: 0.98080003
 BMA accuracy: 0.9866667
 </pre></div>
 </div>

diff --git a/examples/howto_use_pymc.html b/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/607983e7d9ae823d53d48df672ad227d7606464223954292df409a88bf3f1584.png" src="../_images/607983e7d9ae823d53d48df672ad227d7606464223954292df409a88bf3f1584.png" />
+<img alt="../_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png" src="../_images/054a5bfc52a079abe98d328ea4ae2d2437b82acea42eb4a13017a1d7a7130ca9.png" />
 </div>
 </div>
 </section>

diff --git a/examples/howto_use_tfp.html b/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.89
+<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>Average acceptance rate: 0.94
 </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/adc1e7538b93ddc1114cc816750888353642cbd9f3d50883a92c3935f9d93f35.png" src="../_images/adc1e7538b93ddc1114cc816750888353642cbd9f3d50883a92c3935f9d93f35.png" />
+<img alt="../_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png" src="../_images/86e8c57f86009c65e034c8fb58d406ef7907b4f2be4820a6b891352b7a15cd23.png" />
 </div>
 </div>
 </section>