Multi-Path Pathfinder

blackjax/__init__.py

@@ -45,7 +45,7 @@
 from .vi import pathfinder as _pathfinder
 from .vi import schrodinger_follmer as _schrodinger_follmer
 from .vi import svgd as _svgd
-from .vi.pathfinder import PathFinderAlgorithm
+from .vi.pathfinder import MultiPathfinderAlgorithm, multi_pathfinder
 
 """
 The above three classes exist as a backwards compatible way of exposing both the high level, differentiable
@@ -81,10 +81,12 @@ def __call__(self, *args, **kwargs) -> VIAlgorithm:
 @dataclasses.dataclass
 class GeneratePathfinderAPI:
     differentiable: Callable
-    approximate: Callable
-    sample: Callable
+    init: Callable
+    pathfinder: Callable
+    logp: Callable
+    importance_sampling: Callable
 
-    def __call__(self, *args, **kwargs) -> PathFinderAlgorithm:
+    def __call__(self, *args, **kwargs) -> MultiPathfinderAlgorithm:
         return self.differentiable(*args, **kwargs)
 
 
@@ -153,7 +155,11 @@ def generate_top_level_api_from(module):
 )
 
 pathfinder = GeneratePathfinderAPI(
-    _pathfinder.as_top_level_api, _pathfinder.approximate, _pathfinder.sample
+    _pathfinder.as_top_level_api,
+    _pathfinder.init,
+    _pathfinder.pathfinder,
+    _pathfinder.logp,
+    _pathfinder.importance_sampling,
 )
 
 
@@ -169,4 +175,5 @@ def generate_top_level_api_from(module):
     "adjusted_mclmc_find_L_and_step_size",  # adjusted mclmc adaptation
     "ess",  # diagnostics
     "rhat",
+    "multi_pathfinder",  # pathfinder
 ]

blackjax/adaptation/pathfinder_adaptation.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Implementation of the Pathinder warmup for the HMC family of sampling algorithms."""
+
 from typing import Callable, NamedTuple
 
 import jax
@@ -197,7 +198,12 @@ def one_step(carry, rng_key):
             adaptation_info_fn(new_state, info, new_adaptation_state),
         )
 
-    def run(rng_key: PRNGKey, position: ArrayLikeTree, num_steps: int = 400):
+    def run(
+        rng_key: PRNGKey,
+        position: ArrayLikeTree,
+        num_steps: int = 400,
+        num_samples_per_path: int = 1000,
+    ):
         init_key, sample_key, rng_key = jax.random.split(rng_key, 3)
 
         pathfinder_state, _ = vi.pathfinder.approximate(
@@ -210,7 +216,9 @@ def run(rng_key: PRNGKey, position: ArrayLikeTree, num_steps: int = 400):
             initial_step_size,
         )
 
-        init_position, _ = vi.pathfinder.sample(sample_key, pathfinder_state)
+        init_position, _ = vi.pathfinder.sample(
+            sample_key, pathfinder_state, num_samples_per_path
+        )
         init_state = algorithm.init(init_position, logdensity_fn)
 
         keys = jax.random.split(rng_key, num_steps)