fix(pymc): allow data named x with unfrozen model

fixes #157

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -406,8 +406,8 @@ def logp_fn_jax_grad(x, *shared):
         seen.add(val)
 
     def make_logp_func():
-        def logp(x, **shared):
-            logp, grad = logp_fn(x, *[shared[name] for name in logp_shared_names])
+        def logp(_x, **shared):
+            logp, grad = logp_fn(_x, *[shared[name] for name in logp_shared_names])
             return float(logp), np.asarray(grad, dtype="float64", order="C")
 
         return logp
@@ -418,8 +418,8 @@ def logp(x, **shared):
 
     def make_expand_func(seed1, seed2, chain):
         # TODO handle seeds
-        def expand(x, **shared):
-            values = expand_fn(x, *[shared[name] for name in expand_shared_names])
+        def expand(_x, **shared):
+            values = expand_fn(_x, *[shared[name] for name in expand_shared_names])
             return {
                 name: np.asarray(val, order="C", dtype=dtype).ravel()
                 for name, val, dtype in zip(names, values, dtypes, strict=True)

tests/test_pymc.py

@@ -31,6 +31,21 @@ def test_pymc_model(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
+@parameterize_backends
+def test_name_x(backend, gradient_backend):
+    with pm.Model() as model:
+        x = pm.Data("x", 1.0)
+        a = pm.Normal("a", mu=x)
+        pm.Deterministic("z", x * a)
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend, freeze_model=False
+    )
+    trace = nutpie.sample(compiled, chains=1)
+    trace.posterior.a  # noqa: B018
+
+
 @pytest.mark.pymc
 def test_order_shared():
     a_val = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])