fix: Fix jax backend with non-identifier variable names

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -297,8 +297,8 @@ def _compile_pymc_model_jax(model, *, gradient_backend=None, **kwargs):
         orig_logp_fn = logp_fn._fun
 
         @jax.jit
-        def logp_fn_jax_grad(x, **shared):
-            return jax.value_and_grad(lambda x: orig_logp_fn(x, **shared)[0])(x)
+        def logp_fn_jax_grad(x, *shared):
+            return jax.value_and_grad(lambda x: orig_logp_fn(x, *shared)[0])(x)
 
         logp_fn = logp_fn_jax_grad
 
@@ -317,9 +317,7 @@ def logp_fn_jax_grad(x, **shared):
 
     def make_logp_func():
         def logp(x, **shared):
-            logp, grad = logp_fn(
-                x, **{name: shared[name] for name in logp_shared_names}
-            )
+            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
@@ -330,9 +328,7 @@ def logp(x, **shared):
     def make_expand_func(seed1, seed2, chain):
         # TODO handle seeds
         def expand(x, **shared):
-            values = expand_fn(
-                x, **{name: shared[name] for name in expand_shared_names}
-            )
+            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

@@ -127,18 +127,34 @@ def test_det(backend, gradient_backend):
     assert trace.posterior.b.shape[-1] == 2
 
 
+@parameterize_backends
+def test_non_identifier_names(backend, gradient_backend):
+    with pm.Model() as model:
+        a = pm.Uniform("a/b", shape=2)
+        with pm.Model("foo"):
+            c = pm.Data("c", np.array([2.0, 3.0]))
+            pm.Deterministic("b", c * a)
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    trace = nutpie.sample(compiled, chains=1)
+    assert trace.posterior["a/b"].shape[-1] == 2
+    assert trace.posterior["foo::b"].shape[-1] == 2
+
+
 @parameterize_backends
 def test_pymc_model_shared(backend, gradient_backend):
     with pm.Model() as model:
-        mu = pm.Data("mu", 0.1)
+        mu = pm.Data("mu", -0.1)
         sigma = pm.Data("sigma", np.ones(3))
         pm.Normal("a", mu=mu, sigma=sigma, shape=3)
 
     compiled = nutpie.compile_pymc_model(
         model, backend=backend, gradient_backend=gradient_backend
     )
     trace = nutpie.sample(compiled, chains=1, seed=1)
-    np.testing.assert_allclose(trace.posterior.a.mean().values, 0.1, atol=0.05)
+    np.testing.assert_allclose(trace.posterior.a.mean().values, -0.1, atol=0.05)
 
     compiled2 = compiled.with_data(mu=10.0, sigma=3 * np.ones(3))
     trace2 = nutpie.sample(compiled2, chains=1, seed=1)