fix: Move mvscale import to be conditional

pymc-devs · Nov 14, 2024 · 3f64e42 · 3f64e42
1 parent 9736963
commit 3f64e42
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/nutpie/transform_adapter.py b/python/nutpie/transform_adapter.py
@@ -18,7 +18,6 @@ def make_transform_adapter(
     import flowjax
     import flowjax.train
     import flowjax.flows
-    from flowjax.bijections import mvscale
     import optax
     import traceback
     from paramax import Parameterize, unwrap
@@ -164,6 +163,8 @@ def make_layer(key, is_last=False):
                 flow = flowjax.flows._add_default_permute(coupling, n_dim, key_permute)
 
             if scale_layer:
+                from flowjax.bijections import mvscale
+
                 bijections = list(flow.bijections)
                 bijections.append(mvscale.MvScale4(jnp.ones(n_dim) * 1e-5))
                 # bijections.append(mvscale.MvScale3(jnp.ones(n_dim) * 1e-5))