fix: handle missing flowjax correctly

Author: aseyboldt

python/nutpie/compile_stan.py

@@ -1,7 +1,6 @@
 import json
 import tempfile
 from dataclasses import dataclass, replace
-from functools import partial
 from importlib.util import find_spec
 from pathlib import Path
 from typing import Any, Optional
@@ -12,7 +11,6 @@
 
 from nutpie import _lib
 from nutpie.sample import CompiledModel
-from nutpie.transform_adapter import make_transform_adapter
 
 
 class _NumpyArrayEncoder(json.JSONEncoder):
@@ -45,13 +43,14 @@ def with_data(self, *, seed=None, **updates):
         else:
             data_json = None
 
-        kwargs = self._transform_adapt_args
-        if kwargs is None:
-            kwargs = {}
-        make_adapter = partial(
-            make_transform_adapter(**kwargs),
-            logp_fn=None,
-        )
+        outer_kwargs = self._transform_adapt_args
+        if outer_kwargs is None:
+            outer_kwargs = {}
+
+        def make_adapter(*args, **kwargs):
+            from nutpie.transform_adapter import make_transform_adapter
+
+            return make_transform_adapter(**outer_kwargs)(*args, **kwargs, logp_fn=None)
 
         model = _lib.StanModel(self.library, seed, data_json, make_adapter)
         coords = self._coords

python/nutpie/compiled_pyfunc.py

@@ -7,7 +7,6 @@
 
 from nutpie import _lib  # type: ignore
 from nutpie.sample import CompiledModel
-from nutpie.transform_adapter import make_transform_adapter
 
 SeedType = int
 
@@ -67,13 +66,17 @@ def make_expand_func(seed1, seed2, chain):
             return partial(expand_fn, **self._shared_data)
 
         if self._raw_logp_fn is not None:
-            kwargs = self._transform_adapt_args
-            if kwargs is None:
-                kwargs = {}
-            make_adapter = partial(
-                make_transform_adapter(**kwargs),
-                logp_fn=self._raw_logp_fn,
-            )
+            outer_kwargs = self._transform_adapt_args
+            if outer_kwargs is None:
+                outer_kwargs = {}
+
+            def make_adapter(*args, **kwargs):
+                from nutpie.transform_adapter import make_transform_adapter
+
+                return make_transform_adapter(**outer_kwargs)(
+                    *args, **kwargs, logp_fn=self._raw_logp_fn
+                )
+
         else:
             make_adapter = None
 

python/nutpie/transform_adapter.py

@@ -1,7 +1,13 @@
 from functools import partial
+from importlib.util import find_spec
 from typing import Callable
 import time
 
+if find_spec("flowjax") is None:
+    raise ImportError(
+        "The 'flowjax' package is required to use normalizing flow adaptation."
+    )
+
 from flowjax import bijections
 from jaxtyping import ArrayLike, PyTree
 import numpy as np

tests/test_stan.py

@@ -38,3 +38,27 @@ def test_stan_model_data():
         trace = nutpie.sample(compiled_model)
     trace = nutpie.sample(compiled_model.with_data(x=np.array(3.0)))
     trace.posterior.a  # noqa: B018
+
+
+@pytest.mark.slow
+def test_stan_flow():
+    model = """
+    parameters {
+        real a;
+        real<lower=0> b;
+    }
+    model {
+        a ~ normal(0, 1);
+        b ~ normal(0, 1);
+    }
+    """
+
+    compiled_model = nutpie.compile_stan_model(code=model).with_transform_adapt(
+        num_layers=2,
+        nn_width=4,
+        num_diag_windows=6,
+    )
+    trace = nutpie.sample(
+        compiled_model, transform_adapt=True, window_switch_freq=150, tune=600, chains=1
+    )
+    trace.posterior.a  # noqa: B018