Minor changes made to the fit_pathfinder function and added test

`fit_pathfinder`
- Edited `fit_pathfinder` to produce `pathfinder_state`, `pathfinder_info`, `pathfinder_samples` and `pathfinder_idata` for closer examination of the outputs.
- Changed the `num_samples` argument name to `num_draws` to avoid `TypeError` got multiple values for keyword argument 'num_samples'.
- Initial points are automatically set to jitter as jitter is required for pathfinder.

Extras
- New function 'get_jaxified_logp_ravel_inputs' to simplify previous code structure in fit_pathfinder.

Tests
- Added extra test for pathfinder to test pathfinder_info variables and pathfinder_idata  are consistent for a given random seed.

Author: aphc14

pymc_experimental/inference/pathfinder.py

@@ -15,20 +15,55 @@
 import collections
 import sys
 
+from collections.abc import Callable
+
 import arviz as az
 import blackjax
 import jax
 import numpy as np
 import pymc as pm
 
 from packaging import version
+from pymc import Model
 from pymc.backends.arviz import coords_and_dims_for_inferencedata
 from pymc.blocking import DictToArrayBijection, RaveledVars
+from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
+from pymc.model.core import Point
 from pymc.sampling.jax import get_jaxified_graph
 from pymc.util import RandomSeed, _get_seeds_per_chain, get_default_varnames
 
 
+def get_jaxified_logp_ravel_inputs(
+    model: Model,
+    initial_points: dict | None = None,
+) -> tuple[Callable, DictToArrayBijection]:
+    """
+    Get jaxified logp function and ravel inputs for a PyMC model.
+
+    Parameters
+    ----------
+    model : Model
+        PyMC model to jaxify.
+
+    Returns
+    -------
+    tuple[Callable, DictToArrayBijection]
+        A tuple containing the jaxified logp function and the DictToArrayBijection.
+    """
+
+    new_logprob, new_input = pm.pytensorf.join_nonshared_inputs(
+        initial_points, (model.logp(),), model.value_vars, ()
+    )
+
+    logprob_fn_list = get_jaxified_graph([new_input], new_logprob)
+
+    def logprob_fn(x):
+        return logprob_fn_list(x)[0]
+
+    return logprob_fn, DictToArrayBijection.map(initial_points)
+
+
 def convert_flat_trace_to_idata(
     samples,
     include_transformed=False,
@@ -37,7 +72,7 @@ def convert_flat_trace_to_idata(
 ):
     model = modelcontext(model)
     ip = model.initial_point()
-    ip_point_map_info = pm.blocking.DictToArrayBijection.map(ip).point_map_info
+    ip_point_map_info = DictToArrayBijection.map(ip).point_map_info
     trace = collections.defaultdict(list)
     for sample in samples:
         raveld_vars = RaveledVars(sample, ip_point_map_info)
@@ -62,10 +97,10 @@ def convert_flat_trace_to_idata(
 
 
 def fit_pathfinder(
-    num_samples=1000,
+    model=None,
+    num_draws=1000,
     random_seed: RandomSeed | None = None,
     postprocessing_backend="cpu",
-    model=None,
     **pathfinder_kwargs,
 ):
     """
@@ -99,22 +134,19 @@ def fit_pathfinder(
 
     model = modelcontext(model)
 
-    ip = model.initial_point()
-    ip_map = DictToArrayBijection.map(ip)
+    [jitter_seed, pathfinder_seed, sample_seed] = _get_seeds_per_chain(random_seed, 3)
 
-    new_logprob, new_input = pm.pytensorf.join_nonshared_inputs(
-        ip, (model.logp(),), model.value_vars, ()
+    # set initial points. PF requires jittering of initial points
+    ipfn = make_initial_point_fn(
+        model=model,
+        jitter_rvs=set(model.free_RVs),
+        # TODO: add argument for jitter strategy
     )
-
-    logprob_fn_list = get_jaxified_graph([new_input], new_logprob)
-
-    def logprob_fn(x):
-        return logprob_fn_list(x)[0]
-
-    [pathfinder_seed, sample_seed] = _get_seeds_per_chain(random_seed, 2)
+    ip = Point(ipfn(jitter_seed), model=model)
+    logprob_fn, ip_map = get_jaxified_logp_ravel_inputs(model, initial_points=ip)
 
     print("Running pathfinder...", file=sys.stdout)
-    pathfinder_state, _ = blackjax.vi.pathfinder.approximate(
+    pathfinder_state, pathfinder_info = blackjax.vi.pathfinder.approximate(
         rng_key=jax.random.key(pathfinder_seed),
         logdensity_fn=logprob_fn,
         initial_position=ip_map.data,
@@ -123,12 +155,12 @@ def logprob_fn(x):
     pathfinder_samples, _ = blackjax.vi.pathfinder.sample(
         rng_key=jax.random.key(sample_seed),
         state=pathfinder_state,
-        num_samples=num_samples,
+        num_samples=num_draws,
     )
 
     idata = convert_flat_trace_to_idata(
         pathfinder_samples,
         postprocessing_backend=postprocessing_backend,
         model=model,
     )
-    return idata
+    return pathfinder_state, pathfinder_info, pathfinder_samples, idata

tests/test_pathfinder.py

@@ -17,12 +17,14 @@
 import numpy as np
 import pymc as pm
 import pytest
+import xarray as xr
 
 import pymc_experimental as pmx
 
+from pymc_experimental.inference.pathfinder import fit_pathfinder
 
-@pytest.mark.skipif(sys.platform == "win32", reason="JAX not supported on windows.")
-def test_pathfinder():
+
+def build_eight_schools_model():
     # Data of the Eight Schools Model
     J = 8
     y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
@@ -35,6 +37,14 @@ def test_pathfinder():
         theta = pm.Normal("theta", mu=0, sigma=1, shape=J)
         obs = pm.Normal("obs", mu=mu + tau * theta, sigma=sigma, shape=J, observed=y)
 
+    return model
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="JAX not supported on windows.")
+def test_pathfinder():
+    model = build_eight_schools_model()
+
+    with model:
         idata = pmx.fit(method="pathfinder", random_seed=41)
 
     assert idata.posterior["mu"].shape == (1, 1000)
@@ -43,3 +53,65 @@ def test_pathfinder():
     # FIXME: pathfinder doesn't find a reasonable mean! Fix bug or choose model pathfinder can handle
     # np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0)
     np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=0.5)
+
+
+def test_pathfinder_pmx_equivalence():
+    model = build_eight_schools_model()
+    with model:
+        idata_pmx = pmx.fit(method="pathfinder", random_seed=41)
+        idata_pmx = idata_pmx[-1]
+
+    ntests = 2
+    runs = dict()
+    for k in range(ntests):
+        runs[k] = {}
+        (
+            runs[k]["pathfinder_state"],
+            runs[k]["pathfinder_info"],
+            runs[k]["pathfinder_samples"],
+            runs[k]["pathfinder_idata"],
+        ) = fit_pathfinder(model=model, random_seed=41)
+
+        runs[k]["finite_idx"] = (
+            np.argwhere(np.isfinite(runs[k]["pathfinder_info"].path.elbo)).ravel()[-1] + 1
+        )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.elbo[: runs[0]["finite_idx"]],
+        runs[1]["pathfinder_info"].path.elbo[: runs[1]["finite_idx"]],
+    )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.alpha,
+        runs[1]["pathfinder_info"].path.alpha,
+    )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.beta,
+        runs[1]["pathfinder_info"].path.beta,
+    )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.gamma,
+        runs[1]["pathfinder_info"].path.gamma,
+    )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.position,
+        runs[1]["pathfinder_info"].path.position,
+    )
+
+    np.testing.assert_allclose(
+        runs[0]["pathfinder_info"].path.grad_position,
+        runs[1]["pathfinder_info"].path.grad_position,
+    )
+
+    xr.testing.assert_allclose(
+        idata_pmx.posterior,
+        runs[0]["pathfinder_idata"].posterior,
+    )
+
+    xr.testing.assert_allclose(
+        idata_pmx.posterior,
+        runs[1]["pathfinder_idata"].posterior,
+    )