Deploying to gh-pages from @ 3dc3809 🚀

blackjax-devs · Mar 25, 2024 · c822362 · c822362
1 parent 34e8473
commit c822362
Show file tree

Hide file tree

Showing 53 changed files with 188 additions and 96 deletions.
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/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/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/17f157036e3a54e5c33a70b4dce724b28870e82666f77badb7761b1858f4eb43.png b/_images/17f157036e3a54e5c33a70b4dce724b28870e82666f77badb7761b1858f4eb43.png
diff --git a/_images/311b59543904592389a3bb825584938a761aba9e5ff627e97e32285c61c5ad41.png b/_images/311b59543904592389a3bb825584938a761aba9e5ff627e97e32285c61c5ad41.png
diff --git a/_images/329d12c3569a46f98fd06961e8ddf53908ecc8a1fbb8839069a5484ad349081c.png b/_images/329d12c3569a46f98fd06961e8ddf53908ecc8a1fbb8839069a5484ad349081c.png
diff --git a/_images/351927d684a896c23bc7502237f0259cb94c1a01be0612103aa260aef1edcc2f.png b/_images/351927d684a896c23bc7502237f0259cb94c1a01be0612103aa260aef1edcc2f.png
diff --git a/_images/7421d2d2ee38dcb9457c23879d9d5ef68450fd492d9c9ad596c58ab07d62095d.png b/_images/7421d2d2ee38dcb9457c23879d9d5ef68450fd492d9c9ad596c58ab07d62095d.png
diff --git a/_images/7aaefb1bf9ce970ad363b4d21b8d427536cd01cecf66cec871bf753f9b48b3f7.png b/_images/7aaefb1bf9ce970ad363b4d21b8d427536cd01cecf66cec871bf753f9b48b3f7.png
diff --git a/_images/7c1235b7e45c5babac2cb411a000faad8b0eac047423475b3af2f98e0e075367.png b/_images/7c1235b7e45c5babac2cb411a000faad8b0eac047423475b3af2f98e0e075367.png
diff --git a/_images/9b5c946800e2b7cc26ff3b836a0791762a3a31645acb761a05266c1be0197933.png b/_images/9b5c946800e2b7cc26ff3b836a0791762a3a31645acb761a05266c1be0197933.png
diff --git a/_images/9dc13bb0f0e4e38b063ac9586c243abff8ea5fcf54a74fad2797749e66db58f2.png b/_images/9dc13bb0f0e4e38b063ac9586c243abff8ea5fcf54a74fad2797749e66db58f2.png
diff --git a/_images/ab20de987c9499de1352051cdc31e55c5b9a1fd74eed74dc62e10643fea270e4.png b/_images/ab20de987c9499de1352051cdc31e55c5b9a1fd74eed74dc62e10643fea270e4.png
diff --git a/_images/b0bacf20f9b127c76deebb60093bceb2ce856b78ebc18323d333fad341cc0b4a.png b/_images/b0bacf20f9b127c76deebb60093bceb2ce856b78ebc18323d333fad341cc0b4a.png
diff --git a/_images/b121a8653e2f0b21433f76c6c04c81d5316201d73242b412e74a9c84ddbb64f5.png b/_images/b121a8653e2f0b21433f76c6c04c81d5316201d73242b412e74a9c84ddbb64f5.png
diff --git a/_images/bd43216e9dc12f6d927e0f44a0daf9e75b3f8d8799e173b3112246824fa4cce8.png b/_images/bd43216e9dc12f6d927e0f44a0daf9e75b3f8d8799e173b3112246824fa4cce8.png
diff --git a/_images/bf8a1039564e7139dcd6992d792f0db295383f7405ab2150207723789d2c6a95.png b/_images/bf8a1039564e7139dcd6992d792f0db295383f7405ab2150207723789d2c6a95.png
diff --git a/_images/cf2bff6aa74ac128b1bcbf5a93dc10245f5aac3baafdfe445a1c3abcf2bd9ff6.png b/_images/cf2bff6aa74ac128b1bcbf5a93dc10245f5aac3baafdfe445a1c3abcf2bd9ff6.png
diff --git a/_images/db4b95f776e46081ce55027522023543cfd6ef4608d2b97d25580074d85e5953.png b/_images/db4b95f776e46081ce55027522023543cfd6ef4608d2b97d25580074d85e5953.png
diff --git a/_images/e12a7c40431b5f8da0222fefe5c4b341b8e5e5e27163e07fd969cf52b237773f.png b/_images/e12a7c40431b5f8da0222fefe5c4b341b8e5e5e27163e07fd969cf52b237773f.png
diff --git a/_images/f81b5eb0e818b6fb6fc41690f6b4bd22f5bfac13f1569fed2f3b113c22582292.png b/_images/f81b5eb0e818b6fb6fc41690f6b4bd22f5bfac13f1569fed2f3b113c22582292.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+g2ccdfb0&#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;g2ccdfb0&#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+g3dc3809&#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;g3dc3809&#39;</span><span class="p">)</span>
 </pre></div>
 
                 </article>

diff --git a/_modules/blackjax/smc/base.html b/_modules/blackjax/smc/base.html
@@ -415,6 +415,10 @@ <h1>Source code for blackjax.smc.base</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="SMCState.weights">
 <a class="viewcode-back" href="../../../autoapi/blackjax/smc/base/index.html#blackjax.smc.base.SMCState.weights">[docs]</a>
     <span class="n">weights</span><span class="p">:</span> <span class="n">Array</span></div>
+
+<div class="viewcode-block" id="SMCState.update_parameters">
+<a class="viewcode-back" href="../../../autoapi/blackjax/smc/base/index.html#blackjax.smc.base.SMCState.update_parameters">[docs]</a>
+    <span class="n">update_parameters</span><span class="p">:</span> <span class="n">ArrayTree</span></div>
 </div>
 
 
@@ -450,12 +454,12 @@ <h1>Source code for blackjax.smc.base</h1><div class="highlight"><pre>
 
 <div class="viewcode-block" id="init">
 <a class="viewcode-back" href="../../../autoapi/blackjax/smc/base/index.html#blackjax.smc.base.init">[docs]</a>
-<span class="k">def</span> <span class="nf">init</span><span class="p">(</span><span class="n">particles</span><span class="p">:</span> <span class="n">ArrayLikeTree</span><span class="p">):</span>
+<span class="k">def</span> <span class="nf">init</span><span class="p">(</span><span class="n">particles</span><span class="p">:</span> <span class="n">ArrayLikeTree</span><span class="p">,</span> <span class="n">init_update_params</span><span class="p">):</span>
     <span class="c1"># Infer the number of particles from the size of the leading dimension of</span>
     <span class="c1"># the first leaf of the inputted PyTree.</span>
     <span class="n">num_particles</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_util</span><span class="o">.</span><span class="n">tree_flatten</span><span class="p">(</span><span class="n">particles</span><span class="p">)[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
     <span class="n">weights</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">num_particles</span><span class="p">)</span> <span class="o">/</span> <span class="n">num_particles</span>
-    <span class="k">return</span> <span class="n">SMCState</span><span class="p">(</span><span class="n">particles</span><span class="p">,</span> <span class="n">weights</span><span class="p">)</span></div>
+    <span class="k">return</span> <span class="n">SMCState</span><span class="p">(</span><span class="n">particles</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="n">init_update_params</span><span class="p">)</span></div>
 
 
 
@@ -531,17 +535,31 @@ <h1>Source code for blackjax.smc.base</h1><div class="highlight"><pre>
     <span class="n">particles</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">x</span><span class="p">:</span> <span class="n">x</span><span class="p">[</span><span class="n">resampling_idx</span><span class="p">],</span> <span class="n">state</span><span class="o">.</span><span class="n">particles</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">updating_key</span><span class="p">,</span> <span class="n">num_resampled</span><span class="p">)</span>
-    <span class="n">particles</span><span class="p">,</span> <span class="n">update_info</span> <span class="o">=</span> <span class="n">update_fn</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">particles</span><span class="p">)</span>
+    <span class="n">particles</span><span class="p">,</span> <span class="n">update_info</span> <span class="o">=</span> <span class="n">update_fn</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">particles</span><span class="p">,</span> <span class="n">state</span><span class="o">.</span><span class="n">update_parameters</span><span class="p">)</span>
 
     <span class="n">log_weights</span> <span class="o">=</span> <span class="n">weight_fn</span><span class="p">(</span><span class="n">particles</span><span class="p">)</span>
     <span class="n">logsum_weights</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">scipy</span><span class="o">.</span><span class="n">special</span><span class="o">.</span><span class="n">logsumexp</span><span class="p">(</span><span class="n">log_weights</span><span class="p">)</span>
     <span class="n">normalizing_constant</span> <span class="o">=</span> <span class="n">logsum_weights</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">num_particles</span><span class="p">)</span>
     <span class="n">weights</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">log_weights</span> <span class="o">-</span> <span class="n">logsum_weights</span><span class="p">)</span>
 
-    <span class="k">return</span> <span class="n">SMCState</span><span class="p">(</span><span class="n">particles</span><span class="p">,</span> <span class="n">weights</span><span class="p">),</span> <span class="n">SMCInfo</span><span class="p">(</span>
+    <span class="k">return</span> <span class="n">SMCState</span><span class="p">(</span><span class="n">particles</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="n">state</span><span class="o">.</span><span class="n">update_parameters</span><span class="p">),</span> <span class="n">SMCInfo</span><span class="p">(</span>
         <span class="n">resampling_idx</span><span class="p">,</span> <span class="n">normalizing_constant</span><span class="p">,</span> <span class="n">update_info</span>
     <span class="p">)</span></div>
 
+
+
+<div class="viewcode-block" id="extend_params">
+<a class="viewcode-back" href="../../../autoapi/blackjax/smc/base/index.html#blackjax.smc.base.extend_params">[docs]</a>
+<span class="k">def</span> <span class="nf">extend_params</span><span class="p">(</span><span class="n">n_particles</span><span class="p">,</span> <span class="n">params</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Given a dictionary of params, repeats them for every single particle. The expected</span>
+<span class="sd">    usage is in cases where the aim is to repeat the same parameters for all chains within SMC.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="nf">extend</span><span class="p">(</span><span class="n">param</span><span class="p">):</span>
+        <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">param</span><span class="p">)[</span><span class="kc">None</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">n_particles</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="k">return</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span><span class="n">extend</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span></div>
+
 </pre></div>
 
                 </article>