Deploying to gh-pages from @ a5f7482 🚀

blackjax-devs · Apr 8, 2024 · c7977f0 · c7977f0
1 parent a4589bd
commit c7977f0
Show file tree

Hide file tree

Showing 63 changed files with 131 additions and 137 deletions.
diff --git a/.doctrees/autoapi/blackjax/mcmc/trajectory/index.doctree b/.doctrees/autoapi/blackjax/mcmc/trajectory/index.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_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/09b7cbe33887701ad1884aa37df7b2a56b89e22301b128cf6e2addd7f0b65a1a.png b/_images/09b7cbe33887701ad1884aa37df7b2a56b89e22301b128cf6e2addd7f0b65a1a.png
diff --git a/_images/1c5b6fe7b395a5d2e7981da63bd5847ceadd738fc961d8b35ee0821af7d4bd00.png b/_images/1c5b6fe7b395a5d2e7981da63bd5847ceadd738fc961d8b35ee0821af7d4bd00.png
diff --git a/_images/2875ed7733fd1f5cf1ee77ef175b1d1cad67c57616128bc5cb18e869e4c273ba.png b/_images/2875ed7733fd1f5cf1ee77ef175b1d1cad67c57616128bc5cb18e869e4c273ba.png
diff --git a/_images/33d9e22ba4fd431ee4902a212f60c2df78b0c836b16520e104464dd728bfc465.png b/_images/33d9e22ba4fd431ee4902a212f60c2df78b0c836b16520e104464dd728bfc465.png
diff --git a/_images/3a29e9453d2a749b63e822ea31c0a540c6a41e85d5659f8b99a0c9dae9c303ad.png b/_images/3a29e9453d2a749b63e822ea31c0a540c6a41e85d5659f8b99a0c9dae9c303ad.png
diff --git a/_images/5c3b5400ea411e9f5fa2e45c13abe0d8d4e1dc68158a425140f91cf21fb2adbf.png b/_images/5c3b5400ea411e9f5fa2e45c13abe0d8d4e1dc68158a425140f91cf21fb2adbf.png
diff --git a/_images/5ce4826f04084183d974889cd48c63dd6b3c40f14d901f4794f3d5ca981a95fa.png b/_images/5ce4826f04084183d974889cd48c63dd6b3c40f14d901f4794f3d5ca981a95fa.png
diff --git a/_images/5dc9b8e4212d1ab682eb0b0da095f50f819b6bd8c7cea7589a928f6c49c74a99.png b/_images/5dc9b8e4212d1ab682eb0b0da095f50f819b6bd8c7cea7589a928f6c49c74a99.png
diff --git a/_images/6d24c7d128bd1e7aa5b015a6d672a0a4a101cc46b5bd1b5d253ae224e8ae3a0b.png b/_images/6d24c7d128bd1e7aa5b015a6d672a0a4a101cc46b5bd1b5d253ae224e8ae3a0b.png
diff --git a/_images/a31341f653e2d2c92b9691d0ac0d9eeacb82717637b7565b2aff4a0d9dba0782.png b/_images/a31341f653e2d2c92b9691d0ac0d9eeacb82717637b7565b2aff4a0d9dba0782.png
diff --git a/_images/a36d61a6769df0fa65ad90eae4b75b7b01d29cbd5288c48eb36ad7a5bebd790e.png b/_images/a36d61a6769df0fa65ad90eae4b75b7b01d29cbd5288c48eb36ad7a5bebd790e.png
diff --git a/_images/b66d1af8cffedc9bcb8e05fb6b6d2cefcee7080f25e1d6af0c2cb4d95ff6bcfb.png b/_images/b66d1af8cffedc9bcb8e05fb6b6d2cefcee7080f25e1d6af0c2cb4d95ff6bcfb.png
diff --git a/_images/c0437398fe353ed60e11565dcb1f0d5ac0db737f627f388b0cbb15a62dfc7626.png b/_images/c0437398fe353ed60e11565dcb1f0d5ac0db737f627f388b0cbb15a62dfc7626.png
diff --git a/_images/c2f751674dd3e05fa9af64a29c0897464d1072e7d42d1a07e4be50664c012bd6.png b/_images/c2f751674dd3e05fa9af64a29c0897464d1072e7d42d1a07e4be50664c012bd6.png
diff --git a/_images/c7a77283283efc89df6b2a2df84e4c09c0521e01350c373eb262443ee0bc39b7.png b/_images/c7a77283283efc89df6b2a2df84e4c09c0521e01350c373eb262443ee0bc39b7.png
diff --git a/_images/ee43c77154a5f66eddccfed7846804cae4c364db1e6d89c699a9925f557ba4df.png b/_images/ee43c77154a5f66eddccfed7846804cae4c364db1e6d89c699a9925f557ba4df.png
diff --git a/_images/ef06a88c32cae289079c9b0900428a14d72207367250dbbfbfd1fdf9ffb9d565.png b/_images/ef06a88c32cae289079c9b0900428a14d72207367250dbbfbfd1fdf9ffb9d565.png
diff --git a/_images/ef9f3a2e20410436e2ee718043fcbc863b3c7777f550e5b9feb53a15deae8a6c.png b/_images/ef9f3a2e20410436e2ee718043fcbc863b3c7777f550e5b9feb53a15deae8a6c.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+g7cf4f9d&#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;g7cf4f9d&#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+ga5f7482&#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;ga5f7482&#39;</span><span class="p">)</span>
 </pre></div>
 
                 </article>

diff --git a/_modules/blackjax/adaptation/chees_adaptation.html b/_modules/blackjax/adaptation/chees_adaptation.html
@@ -764,20 +764,18 @@ <h1>Source code for blackjax.adaptation.chees_adaptation</h1><div class="highlig
         <span class="p">),</span> <span class="s2">&quot;initial `positions` leading dimension must be equal to the `num_chains`&quot;</span>
         <span class="n">num_dim</span> <span class="o">=</span> <span class="n">pytree_size</span><span class="p">(</span><span class="n">positions</span><span class="p">)</span> <span class="o">//</span> <span class="n">num_chains</span>
 
-        <span class="n">key_init</span><span class="p">,</span> <span class="n">key_step</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">next_random_arg_fn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span>
+        <span class="n">init_random_arg</span> <span class="o">=</span> <span class="mi">0</span>
 
         <span class="k">if</span> <span class="n">jitter_generator</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
-            <span class="n">jitter_gn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">key</span><span class="p">:</span> <span class="n">jitter_generator</span><span class="p">(</span><span class="n">key</span><span class="p">)</span> <span class="o">*</span> <span class="n">jitter_amount</span> <span class="o">+</span> <span class="p">(</span>
-                <span class="mf">1.0</span> <span class="o">-</span> <span class="n">jitter_amount</span>
-            <span class="p">)</span>
-            <span class="n">next_random_arg_fn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">key</span><span class="p">:</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">key</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span>
-            <span class="n">init_random_arg</span> <span class="o">=</span> <span class="n">key_init</span>
+            <span class="n">rng_key</span><span class="p">,</span> <span class="n">carry_key</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">jitter_gn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">jitter_generator</span><span class="p">(</span>
+                <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">fold_in</span><span class="p">(</span><span class="n">carry_key</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span>
+            <span class="p">)</span> <span class="o">*</span> <span class="n">jitter_amount</span> <span class="o">+</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">jitter_amount</span><span class="p">)</span>
         <span class="k">else</span><span class="p">:</span>
             <span class="n">jitter_gn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">dynamic_hmc</span><span class="o">.</span><span class="n">halton_sequence</span><span class="p">(</span>
                 <span class="n">i</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">log2</span><span class="p">(</span><span class="n">num_steps</span> <span class="o">+</span> <span class="n">max_sampling_steps</span><span class="p">))</span>
             <span class="p">)</span> <span class="o">*</span> <span class="n">jitter_amount</span> <span class="o">+</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">jitter_amount</span><span class="p">)</span>
-            <span class="n">next_random_arg_fn</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span>
-            <span class="n">init_random_arg</span> <span class="o">=</span> <span class="mi">0</span>
 
         <span class="k">def</span> <span class="nf">integration_steps_fn</span><span class="p">(</span><span class="n">random_generator_arg</span><span class="p">,</span> <span class="n">trajectory_length_adjusted</span><span class="p">):</span>
             <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span>
@@ -828,7 +826,7 @@ <h1>Source code for blackjax.adaptation.chees_adaptation</h1><div class="highlig
         <span class="n">init_states</span> <span class="o">=</span> <span class="n">batch_init</span><span class="p">(</span><span class="n">positions</span><span class="p">)</span>
         <span class="n">init_adaptation_state</span> <span class="o">=</span> <span class="n">init</span><span class="p">(</span><span class="n">init_random_arg</span><span class="p">,</span> <span class="n">step_size</span><span class="p">)</span>
 
-        <span class="n">keys_step</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">key_step</span><span class="p">,</span> <span class="n">num_steps</span><span class="p">)</span>
+        <span class="n">keys_step</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="p">(</span><span class="n">last_states</span><span class="p">,</span> <span class="n">last_adaptation_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="p">(</span><span class="n">init_states</span><span class="p">,</span> <span class="n">init_adaptation_state</span><span class="p">),</span> <span class="n">keys_step</span>
         <span class="p">)</span>

diff --git a/_modules/blackjax/adaptation/mclmc_adaptation.html b/_modules/blackjax/adaptation/mclmc_adaptation.html
@@ -613,12 +613,12 @@ <h1>Source code for blackjax.adaptation.mclmc_adaptation</h1><div class="highlig
         <span class="n">kalman_state</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">dim</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">dim</span><span class="p">))</span>
 
         <span class="c1"># run the steps</span>
-        <span class="n">kalman_state</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">kalman_state</span><span class="p">,</span> <span class="o">*</span><span class="n">_</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">step</span><span class="p">,</span>
             <span class="n">init</span><span class="o">=</span><span class="p">(</span><span class="n">state</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">adap0</span><span class="p">,</span> <span class="n">kalman_state</span><span class="p">),</span>
             <span class="n">xs</span><span class="o">=</span><span class="p">(</span><span class="n">outer_weights</span><span class="p">,</span> <span class="n">L_step_size_adaptation_keys</span><span class="p">),</span>
             <span class="n">length</span><span class="o">=</span><span class="n">num_steps1</span> <span class="o">+</span> <span class="n">num_steps2</span><span class="p">,</span>
-        <span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
+        <span class="p">)</span>
         <span class="n">state</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">kalman_state_output</span> <span class="o">=</span> <span class="n">kalman_state</span>
 
         <span class="n">L</span> <span class="o">=</span> <span class="n">params</span><span class="o">.</span><span class="n">L</span>

diff --git a/_modules/blackjax/adaptation/meads_adaptation.html b/_modules/blackjax/adaptation/meads_adaptation.html
@@ -481,16 +481,16 @@ <h1>Source code for blackjax.adaptation.meads_adaptation</h1><div class="highlig
 <span class="sd">        of the generalized HMC algorithm.</span>
 
 <span class="sd">        &quot;&quot;&quot;</span>
-        <span class="n">mean_position</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">p</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">),</span> <span class="n">positions</span><span class="p">)</span>
-        <span class="n">sd_position</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">p</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">),</span> <span class="n">positions</span><span class="p">)</span>
-        <span class="n">normalized_positions</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span>
+        <span class="n">mean_position</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">p</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">),</span> <span class="n">positions</span><span class="p">)</span>
+        <span class="n">sd_position</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">p</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">),</span> <span class="n">positions</span><span class="p">)</span>
+        <span class="n">normalized_positions</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree</span><span class="o">.</span><span class="n">map</span><span class="p">(</span>
             <span class="k">lambda</span> <span class="n">p</span><span class="p">,</span> <span class="n">mu</span><span class="p">,</span> <span class="n">sd</span><span class="p">:</span> <span class="p">(</span><span class="n">p</span> <span class="o">-</span> <span class="n">mu</span><span class="p">)</span> <span class="o">/</span> <span class="n">sd</span><span class="p">,</span>
             <span class="n">positions</span><span class="p">,</span>
             <span class="n">mean_position</span><span class="p">,</span>
             <span class="n">sd_position</span><span class="p">,</span>
         <span class="p">)</span>
 
-        <span class="n">batch_grad_scaled</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span>
+        <span class="n">batch_grad_scaled</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree</span><span class="o">.</span><span class="n">map</span><span class="p">(</span>
             <span class="k">lambda</span> <span class="n">grad</span><span class="p">,</span> <span class="n">sd</span><span class="p">:</span> <span class="n">grad</span> <span class="o">*</span> <span class="n">sd</span><span class="p">,</span> <span class="n">logdensity_grad</span><span class="p">,</span> <span class="n">sd_position</span>
         <span class="p">)</span>