pymc-extras 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

pymc_extras/model/model_api.py

@@ -1,6 +1,9 @@
 from functools import wraps
+from inspect import signature
 
-from pymc import Model
+import pytensor.tensor as pt
+
+from pymc import Data, Model
 
 
 def as_model(*model_args, **model_kwargs):
@@ -9,6 +12,8 @@ def as_model(*model_args, **model_kwargs):
     This removes all need to think about context managers and lets you separate creating a generative model from using the model.
     Additionally, a coords argument is added to the function so coords can be changed during function invocation
 
+    All parameters are wrapped with a `pm.Data` object if the underlying type of the data supports it.
+
     Adapted from `Rob Zinkov's blog post <https://www.zinkov.com/posts/2023-alternative-frontends-pymc/>`_ and inspired by the `sampled <https://github.com/colcarroll/sampled>`_ decorator for PyMC3.
 
     Examples
@@ -47,8 +52,19 @@ def as_model(*model_args, **model_kwargs):
         @wraps(f)
         def make_model(*args, **kwargs):
             coords = model_kwargs.pop("coords", {}) | kwargs.pop("coords", {})
+            sig = signature(f)
+            ba = sig.bind(*args, **kwargs)
+            ba.apply_defaults()
+
             with Model(*model_args, coords=coords, **model_kwargs) as m:
-                f(*args, **kwargs)
+                for name, v in ba.arguments.items():
+                    # Only wrap pm.Data around values pytensor can process
+                    try:
+                        _ = pt.as_tensor_variable(v)
+                        ba.arguments[name] = Data(name, v)
+                    except (NotImplementedError, TypeError, ValueError):
+                        pass
+                f(*ba.args, **ba.kwargs)
             return m
 
         return make_model

pymc_extras/statespace/core/statespace.py

@@ -15,6 +15,9 @@ from pymc.model.transform.optimization import freeze_dims_and_data
 from pymc.util import RandomState
 from pytensor import Variable, graph_replace
 from pytensor.compile import get_mode
+from rich.box import SIMPLE_HEAD
+from rich.console import Console
+from rich.table import Table
 
 from pymc_extras.statespace.core.representation import PytensorRepresentation
 from pymc_extras.statespace.filters import (
@@ -254,53 +257,72 @@ class PyMCStateSpace:
         self.kalman_smoother = KalmanSmoother()
         self.make_symbolic_graph()
 
-        if verbose:
-            # These are split into separate try-except blocks, because it will be quite rare of models to implement
-            # _print_data_requirements, but we still want to print the prior requirements.
-            try:
-                self._print_prior_requirements()
-            except NotImplementedError:
-                pass
-            try:
-                self._print_data_requirements()
-            except NotImplementedError:
-                pass
-
-    def _print_prior_requirements(self) -> None:
-        """
-        Prints a short report to the terminal about the priors needed for the model, including their names,
+        self.requirement_table = None
+        self._populate_prior_requirements()
+        self._populate_data_requirements()
+
+        if verbose and self.requirement_table:
+            console = Console()
+            console.print(self.requirement_table)
+
+    def _populate_prior_requirements(self) -> None:
+        """
+        Add requirements about priors needed for the model to a rich table, including their names,
         shapes, named dimensions, and any parameter constraints.
         """
-        out = ""
-        for param, info in self.param_info.items():
-            out += f'\t{param} -- shape: {info["shape"]}, constraints: {info["constraints"]}, dims: {info["dims"]}\n'
-        out = out.rstrip()
+        # Check that the param_info class is implemented, and also that it's a dictionary. We can't proceed if either
+        # is not true.
+        try:
+            if not isinstance(self.param_info, dict):
+                return
+        except NotImplementedError:
+            return
 
-        _log.info(
-            "The following parameters should be assigned priors inside a PyMC "
-            f"model block: \n"
-            f"{out}"
-        )
+        if self.requirement_table is None:
+            self._initialize_requirement_table()
+
+        for param, info in self.param_info.items():
+            self.requirement_table.add_row(
+                param, str(info["shape"]), info["constraints"], str(info["dims"])
+            )
 
-    def _print_data_requirements(self) -> None:
+    def _populate_data_requirements(self) -> None:
         """
-        Prints a short report to the terminal about the data needed for the model, including their names, shapes,
-        and named dimensions.
+        Add requirements about the data needed for the model, including their names, shapes, and named dimensions.
         """
-        if not self.data_info:
+        try:
+            if not isinstance(self.data_info, dict):
+                return
+        except NotImplementedError:
             return
 
-        out = ""
+        if self.requirement_table is None:
+            self._initialize_requirement_table()
+        else:
+            self.requirement_table.add_section()
+
         for data, info in self.data_info.items():
-            out += f'\t{data} -- shape: {info["shape"]}, dims: {info["dims"]}\n'
-        out = out.rstrip()
+            self.requirement_table.add_row(data, str(info["shape"]), "pm.Data", str(info["dims"]))
+
+    def _initialize_requirement_table(self) -> None:
+        self.requirement_table = Table(
+            show_header=True,
+            show_edge=True,
+            box=SIMPLE_HEAD,
+            highlight=True,
+        )
 
-        _log.info(
-            "The following Data variables should be assigned to the model inside a PyMC "
-            f"model block: \n"
-            f"{out}"
+        self.requirement_table.title = "Model Requirements"
+        self.requirement_table.caption = (
+            "These parameters should be assigned priors inside a PyMC model block before "
+            "calling the build_statespace_graph method."
         )
 
+        self.requirement_table.add_column("Variable", justify="left")
+        self.requirement_table.add_column("Shape", justify="left")
+        self.requirement_table.add_column("Constraints", justify="left")
+        self.requirement_table.add_column("Dimensions", justify="right")
+
     def _unpack_statespace_with_placeholders(
         self,
     ) -> tuple[
@@ -961,10 +983,31 @@ class PyMCStateSpace:
         list[pm.Flat]
             A list of pm.Flat variables representing all parameters estimated by the model.
         """
+
+        def infer_variable_shape(name):
+            shape = self._name_to_variable[name].type.shape
+            if not any(dim is None for dim in shape):
+                return shape
+
+            dim_names = self._fit_dims.get(name, None)
+            if dim_names is None:
+                raise ValueError(
+                    f"Could not infer shape for {name}, because it was not given coords during model"
+                    f"fitting"
+                )
+
+            shape_from_coords = tuple([len(self._fit_coords[dim]) for dim in dim_names])
+            return tuple(
+                [
+                    shape[i] if shape[i] is not None else shape_from_coords[i]
+                    for i in range(len(shape))
+                ]
+            )
+
         for name in self.param_names:
             pm.Flat(
                 name,
-                shape=self._name_to_variable[name].type.shape,
+                shape=infer_variable_shape(name),
                 dims=self._fit_dims.get(name, None),
             )
 

pymc_extras/version.txt

@@ -1 +1 @@
-0.2.1
+0.2.2

{pymc_extras-0.2.1.dist-info → pymc_extras-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: pymc-extras
-Version: 0.2.1
+Version: 0.2.2
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Home-page: http://github.com/pymc-devs/pymc-extras
 Maintainer: PyMC Developers
@@ -34,6 +34,17 @@ Provides-Extra: dev
 Requires-Dist: dask[all]; extra == "dev"
 Requires-Dist: blackjax; extra == "dev"
 Requires-Dist: statsmodels; extra == "dev"
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: maintainer
+Dynamic: maintainer-email
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # Welcome to `pymc-extras`
 <a href="https://gitpod.io/#https://github.com/pymc-devs/pymc-extras">

{pymc_extras-0.2.1.dist-info → pymc_extras-0.2.2.dist-info}/RECORD

@@ -3,7 +3,7 @@ pymc_extras/linearmodel.py,sha256=6eitl15Ec15mSZu7zoHZ7Wwy4U1DPwqfAgwEt6ILeIc,39
 pymc_extras/model_builder.py,sha256=sAw77fxdiy046BvDPjocuMlbJ0Efj-CDAGtmcwYmoG0,26361
 pymc_extras/printing.py,sha256=G8mj9dRd6i0PcsbcEWZm56ek6V8mmil78RI4MUhywBs,6506
 pymc_extras/version.py,sha256=VxPGCBzhtSegu-Jp5cjzn0n4DGU0wuPUh-KyZKB6uPM,240
-pymc_extras/version.txt,sha256=cQFcl5zLD8igvnygroMEarBFzcLI-qCfsvD35ED5tKY,6
+pymc_extras/version.txt,sha256=mY9riH7Xpu9E6EIQ0CN7cVvtQupyVPSNxBiv7ApIQQM,6
 pymc_extras/distributions/__init__.py,sha256=gTX7tvX8NcgP7V72URV7GeqF1aAEjGVbuW8LMxhXceY,1295
 pymc_extras/distributions/continuous.py,sha256=z-nvQgGncYISdRY8cWsa-56V0bQGq70jYwU-i8VZ0Uk,11253
 pymc_extras/distributions/discrete.py,sha256=vrARNuiQAEXrs7yQgImV1PO8AV1uyEC_LBhr6F9IcOg,13032
@@ -15,13 +15,16 @@ pymc_extras/gp/__init__.py,sha256=sFHw2y3lEl5tG_FDQHZUonQ_k0DF1JRf0Rp8dpHmge0,74
 pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQDf4,6682
 pymc_extras/inference/__init__.py,sha256=5cXpaQQnW0mJJ3x8wSxmYu63l--Xab5D_gMtjA6Q3uU,666
 pymc_extras/inference/find_map.py,sha256=T0uO8prUI5aBNuR1AN8fbA4cHmLRQLXznwJrfxfe7CA,15723
-pymc_extras/inference/fit.py,sha256=NFEpUaYLJAmDRP1WIPymgnEcXUofkoURYHbEdiTivzQ,1313
+pymc_extras/inference/fit.py,sha256=S9R48dh74s6K0MC9Iys4NAwVjP6rVRfx6SF-kPiR70E,1165
 pymc_extras/inference/laplace.py,sha256=OglOvnxfHLe0VXxBC1-ddVzADR9zgGxUPScM6P6FYo8,21163
-pymc_extras/inference/pathfinder.py,sha256=cmzR2OZCfkdTipT-8pmLuF-MHmLzxotsYlezOWBUM4U,4171
+pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
+pymc_extras/inference/pathfinder/importance_sampling.py,sha256=VvmuaE3aw_Mo3tMwswfF0rqe19mnhOCpzIScaJzjA1Y,6159
+pymc_extras/inference/pathfinder/lbfgs.py,sha256=P0UIOVtspdLzDU6alK-y91qzVAzXjYAXPuGmZ1nRqMo,5715
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=fomZ5voVcWxvhWpeIZV7IHGIJCasT1g0ivC4dC3-0GM,63694
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
 pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/model/model_api.py,sha256=_r6rYQG1tt9Z95QU-jVyHqZ1rs-u7sFMO5HJ5unDV5A,1750
+pymc_extras/model/model_api.py,sha256=UHMfQXxWBujeSiUySU0fDUC5Sd_BjT8FoVz3iBxQH_4,2400
 pymc_extras/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/marginal/distributions.py,sha256=iM1yT7_BmivgUSloQPKE2QXGPgjvLqDMY_OTBGsdAWg,15563
 pymc_extras/model/marginal/graph_analysis.py,sha256=0hWUH_PjfpgneQ3NaT__pWHS1fh50zNbI86kH4Nub0E,15693
@@ -34,7 +37,7 @@ pymc_extras/statespace/__init__.py,sha256=0MtZj7yT6jcyERvITnn-nkhyY8fO6Za4_vV53C
 pymc_extras/statespace/core/__init__.py,sha256=huHEiXAm8zV2MZyZ8GBHp6q7_fnWqveM7lC6ilpb3iE,309
 pymc_extras/statespace/core/compile.py,sha256=1c8Q9D9zeUe7F0z7CH6q1C6ZuLg2_imgk8RoE_KMaFI,1608
 pymc_extras/statespace/core/representation.py,sha256=DwNIun6wdeEA20oWBx5M4govyWTf5JI87aGQ_E6Mb4U,18956
-pymc_extras/statespace/core/statespace.py,sha256=ZElRm9wJvIGG4Pw-3qiQpBkHXRDqS6pfRyuGrBBcZ2Y,95270
+pymc_extras/statespace/core/statespace.py,sha256=K_WVnWKlI6sR2kgriq9sctQVvwXCeAirm14TthDpmRM,96860
 pymc_extras/statespace/filters/__init__.py,sha256=N9Q4D0gAq_ZtT-GtrqiX1HkSg6Orv7o1TbrWUtnbTJE,420
 pymc_extras/statespace/filters/distributions.py,sha256=-9j__vRqL5hKyYFnQr5HKHA5kEFzwiuSccH4mslTOuQ,12900
 pymc_extras/statespace/filters/kalman_filter.py,sha256=HELC3aK4k8EdWlUAk5_F7y7YkIz-Xi_0j2AwRgAXgcc,31949
@@ -63,7 +66,7 @@ tests/test_histogram_approximation.py,sha256=w-xb2Rr0Qft6sm6F3BTmXXnpuqyefC1SUL6
 tests/test_laplace.py,sha256=5ioEyP6AzmMszrtQRz0KWTsCCU35SEhSOdBcYfYzptE,8228
 tests/test_linearmodel.py,sha256=iB8ApNqIX9_nUHoo-Tm51xuPdrva5t4VLLut6qXB5Ao,6906
 tests/test_model_builder.py,sha256=QiINEihBR9rx8xM4Nqlg4urZKoyo58aTKDtxl9SJF1s,11249
-tests/test_pathfinder.py,sha256=FBm0ge6rje5jz9_10h_247E70aKCpkbu1jmzrR7Ar8A,1726
+tests/test_pathfinder.py,sha256=GnSbZJ9QuFW9UVbkWaVgMVqQZTCttOyz_rSflxhQ-EA,4955
 tests/test_pivoted_cholesky.py,sha256=PuMdMSCzO4KdQWpUF4SEBeuH_qsINCIH8TYtmmJ1NKo,692
 tests/test_printing.py,sha256=HnvwwjrjBuxXFAJdyU0K_lvKGLgh4nzHAnhsIUpenbY,5211
 tests/test_prior_from_trace.py,sha256=HOzR3l98pl7TEJquo_kSugED4wBTgHo4-8lgnpmacs8,5516
@@ -75,7 +78,7 @@ tests/distributions/test_discrete.py,sha256=CjjaUpppsvQ6zLzV15ZsbwNOKrDmEdz4VWcl
 tests/distributions/test_discrete_markov_chain.py,sha256=8RCHZXSB8IWjniuKaGGlM_iTWGmdrcOqginxmrAeEJg,9212
 tests/distributions/test_multivariate.py,sha256=LBvBuoT_3rzi8rR38b8L441Y-9Ff0cIXeRBKiEn6kjs,10452
 tests/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/model/test_model_api.py,sha256=SiOMA1NpyQKJ7stYI1ms8ksDPU81lVo8wS8hbqiik-U,776
+tests/model/test_model_api.py,sha256=FJvMTmexovRELZOUcUyk-6Vwk9qSiH7hIFoiArgl5mk,1040
 tests/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/model/marginal/test_distributions.py,sha256=p5f73g4ogxYkdZaBndZV_1ra8TCppXiRlUpaaTwEe-M,5195
 tests/model/marginal/test_graph_analysis.py,sha256=raoj41NusMOj1zzPCrxrlQODqX6Ey8Ft_o32pNTe5qg,6712
@@ -88,15 +91,15 @@ tests/statespace/test_coord_assignment.py,sha256=2GBm46-0eI4QNh4bvp3D7az58stcA5Z
 tests/statespace/test_distributions.py,sha256=WQ_ROyd-PL3cimXTyEtyVaMEVtS7Hue2Z0lN7UnGDyo,9122
 tests/statespace/test_kalman_filter.py,sha256=s2n62FzXl9elU_uqaMNaEaexUfq3SXe3_YvQ2lM6hiQ,11600
 tests/statespace/test_representation.py,sha256=1KAJY4ZaVhb1WdAJLx2UYSXuVYsMNWX98gEDF7P0B4s,6210
-tests/statespace/test_statespace.py,sha256=8ZLLQaxlP5UEJnIMYyIzzAODCxMxs6E5I1hLu2HCdqo,28866
+tests/statespace/test_statespace.py,sha256=JoupFFpG8PmpB_NFV471IuTmyXhEd6_vOISwVCRrBBM,30570
 tests/statespace/test_statespace_JAX.py,sha256=hZOc6xxYdVeATPCKmcHMLOVcuvdzGRzgQQ4RrDenwk8,5279
 tests/statespace/test_structural.py,sha256=HD8OaGbjuH4y3xv_uG-R1xLZpPpcb4-3dbcTeb_imLY,29306
 tests/statespace/utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/statespace/utilities/shared_fixtures.py,sha256=SNw8Bvj1Yw11TxAW6n20Bq0B8oaYtVTiFFEVNH_wnp4,164
 tests/statespace/utilities/statsmodel_local_level.py,sha256=SQAzaYaSDwiVhUQ1iWjt4MgfAd54RuzVtnslIs3xdS8,1225
 tests/statespace/utilities/test_helpers.py,sha256=oH24a6Q45NFFFI3Kx9mhKbxsCvo9ErCorKFoTjDB3-4,9159
-pymc_extras-0.2.1.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.2.1.dist-info/METADATA,sha256=pT1MOjFxsX6lc0q_D3J-2jNW6UaRkCq_0kJemgG4DGU,4894
-pymc_extras-0.2.1.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-pymc_extras-0.2.1.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
-pymc_extras-0.2.1.dist-info/RECORD,,
+pymc_extras-0.2.2.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.2.2.dist-info/METADATA,sha256=9k60kKKNzr7E24gACpTBNOWj-tRTNOiujSZOFD89G5c,5140
+pymc_extras-0.2.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+pymc_extras-0.2.2.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
+pymc_extras-0.2.2.dist-info/RECORD,,

{pymc_extras-0.2.1.dist-info → pymc_extras-0.2.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.7.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/model/test_model_api.py

@@ -25,5 +25,14 @@ def test_logp():
 
     mw2 = model_wrapped2(coords=coords)
 
+    @pmx.as_model()
+    def model_wrapped3(mu):
+        pm.Normal("x", mu, 1.0, dims="obs")
+
+    mw3 = model_wrapped3(0.0, coords=coords)
+    mw4 = model_wrapped3(np.array([np.nan]), coords=coords)
+
     np.testing.assert_equal(model.point_logps(), mw.point_logps())
     np.testing.assert_equal(mw.point_logps(), mw2.point_logps())
+    assert mw3["mu"] in mw3.data_vars
+    assert "mu" not in mw4

tests/statespace/test_statespace.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from functools import partial
 
 import numpy as np
@@ -349,6 +350,59 @@ def test_sampling_methods(group, kind, ss_mod, idata, rng):
             assert not np.any(np.isnan(test_idata[f"{group}_{output}"].values))
 
 
+@pytest.mark.filterwarnings("ignore:Provided data contains missing values")
+def test_sample_conditional_with_time_varying():
+    class TVCovariance(PyMCStateSpace):
+        def __init__(self):
+            super().__init__(k_states=1, k_endog=1, k_posdef=1)
+
+        def make_symbolic_graph(self) -> None:
+            self.ssm["transition", 0, 0] = 1.0
+
+            self.ssm["design", 0, 0] = 1.0
+
+            sigma_cov = self.make_and_register_variable("sigma_cov", (None,))
+            self.ssm["state_cov"] = sigma_cov[:, None, None] ** 2
+
+        @property
+        def param_names(self) -> list[str]:
+            return ["sigma_cov"]
+
+        @property
+        def coords(self) -> dict[str, Sequence[str]]:
+            return make_default_coords(self)
+
+        @property
+        def state_names(self) -> list[str]:
+            return ["level"]
+
+        @property
+        def observed_states(self) -> list[str]:
+            return ["level"]
+
+        @property
+        def shock_names(self) -> list[str]:
+            return ["level"]
+
+    ss_mod = TVCovariance()
+    empty_data = pd.DataFrame(
+        np.nan, index=pd.date_range("2020-01-01", periods=100, freq="D"), columns=["data"]
+    )
+
+    coords = ss_mod.coords
+    coords["time"] = empty_data.index
+    with pm.Model(coords=coords) as mod:
+        log_sigma_cov = pm.Normal("log_sigma_cov", mu=0, sigma=0.1, dims=["time"])
+        pm.Deterministic("sigma_cov", pm.math.exp(log_sigma_cov.cumsum()), dims=["time"])
+
+        ss_mod.build_statespace_graph(data=empty_data)
+
+        prior = pm.sample_prior_predictive(10)
+
+    ss_mod.sample_unconditional_prior(prior)
+    ss_mod.sample_conditional_prior(prior)
+
+
 def _make_time_idx(mod, use_datetime_index=True):
     if use_datetime_index:
         mod._fit_coords["time"] = nile.index

tests/test_pathfinder.py

@@ -18,12 +18,12 @@ import numpy as np
 import pymc as pm
 import pytest
 
+pytestmark = pytest.mark.filterwarnings("ignore:compile_pymc was renamed to compile:FutureWarning")
+
 import pymc_extras as pmx
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="JAX not supported on windows.")
-def test_pathfinder():
-    # Data of the Eight Schools Model
+def eight_schools_model() -> pm.Model:
     J = 8
     y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
     sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
@@ -35,11 +35,139 @@ 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)
 
-        idata = pmx.fit(method="pathfinder", random_seed=41)
+    return model
+
+
+@pytest.fixture
+def reference_idata():
+    model = eight_schools_model()
+    with model:
+        idata = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+        )
+    return idata
+
+
+@pytest.mark.parametrize("inference_backend", ["pymc", "blackjax"])
+def test_pathfinder(inference_backend, reference_idata):
+    if inference_backend == "blackjax" and sys.platform == "win32":
+        pytest.skip("JAX not supported on windows")
+
+    if inference_backend == "blackjax":
+        model = eight_schools_model()
+        with model:
+            idata = pmx.fit(
+                method="pathfinder",
+                num_paths=50,
+                jitter=10.0,
+                random_seed=41,
+                inference_backend=inference_backend,
+            )
+    else:
+        idata = reference_idata
+        np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0, atol=1.6)
+        np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=1.5)
 
     assert idata.posterior["mu"].shape == (1, 1000)
     assert idata.posterior["tau"].shape == (1, 1000)
     assert idata.posterior["theta"].shape == (1, 1000, 8)
-    # 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)
+
+
+@pytest.mark.parametrize("concurrent", ["thread", "process"])
+def test_concurrent_results(reference_idata, concurrent):
+    model = eight_schools_model()
+    with model:
+        idata_conc = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+            concurrent=concurrent,
+        )
+
+    np.testing.assert_allclose(
+        reference_idata.posterior.mu.data.mean(),
+        idata_conc.posterior.mu.data.mean(),
+        atol=0.4,
+    )
+
+    np.testing.assert_allclose(
+        reference_idata.posterior.tau.data.mean(),
+        idata_conc.posterior.tau.data.mean(),
+        atol=0.4,
+    )
+
+
+def test_seed(reference_idata):
+    model = eight_schools_model()
+    with model:
+        idata_41 = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+        )
+
+        idata_123 = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=123,
+            inference_backend="pymc",
+        )
+
+    assert not np.allclose(idata_41.posterior.mu.data.mean(), idata_123.posterior.mu.data.mean())
+
+    assert np.allclose(idata_41.posterior.mu.data.mean(), idata_41.posterior.mu.data.mean())
+
+
+def test_bfgs_sample():
+    import pytensor.tensor as pt
+
+    from pymc_extras.inference.pathfinder.pathfinder import (
+        alpha_recover,
+        bfgs_sample,
+        inverse_hessian_factors,
+    )
+
+    """test BFGS sampling"""
+    Lp1, N = 8, 10
+    L = Lp1 - 1
+    J = 6
+    num_samples = 1000
+
+    # mock data
+    x_data = np.random.randn(Lp1, N)
+    g_data = np.random.randn(Lp1, N)
+
+    # get factors
+    x_full = pt.as_tensor(x_data, dtype="float64")
+    g_full = pt.as_tensor(g_data, dtype="float64")
+    epsilon = 1e-11
+
+    x = x_full[1:]
+    g = g_full[1:]
+    alpha, S, Z, update_mask = alpha_recover(x_full, g_full, epsilon)
+    beta, gamma = inverse_hessian_factors(alpha, S, Z, update_mask, J)
+
+    # sample
+    phi, logq = bfgs_sample(
+        num_samples=num_samples,
+        x=x,
+        g=g,
+        alpha=alpha,
+        beta=beta,
+        gamma=gamma,
+    )
+
+    # check shapes
+    assert beta.eval().shape == (L, N, 2 * J)
+    assert gamma.eval().shape == (L, 2 * J, 2 * J)
+    assert phi.eval().shape == (L, num_samples, N)
+    assert logq.eval().shape == (L, num_samples)