fix: pymc with jax backend was broken with some shapes

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -445,8 +445,10 @@ def make_expand_func(seed1, seed2, chain):
         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)
+                name: np.asarray(val, order="C", dtype=dtype).reshape(shape)
+                for name, val, dtype, shape in zip(
+                    names, values, dtypes, shapes, strict=True
+                )
             }
 
         return expand

src/pyfunc.rs

@@ -289,17 +289,27 @@ impl CpuLogpFunc for PyDensity {
                         ))
                     })?;
                     if !arr.is_c_contiguous() {
-                        return Err(nuts_rs::CpuMathError::ExpandError(
-                            "not c contiguous".into(),
-                        ));
+                        return Err(nuts_rs::CpuMathError::ExpandError(format!(
+                            "not c contiguous: {}",
+                            var.name
+                        )));
+                    }
+                    if arr.shape().len() != var.shape.as_slice().len() {
+                        return Err(nuts_rs::CpuMathError::ExpandError(format!(
+                            "unexpected number of dimensions for variable {}",
+                            var.name
+                        )));
                     }
                     if !arr
                         .shape()
                         .iter()
                         .zip(var.shape.as_slice())
                         .all(|(a, &b)| *a as u64 == b)
                     {
-                        return Err(nuts_rs::CpuMathError::ExpandError("upected shape".into()));
+                        return Err(nuts_rs::CpuMathError::ExpandError(format!(
+                            "unexpected shape for variable {}",
+                            var.name
+                        )));
                     }
                     Ok(arr)
                 }

tests/test_pymc.py

@@ -100,7 +100,9 @@ def test_low_rank(backend, gradient_backend):
 @parameterize_backends
 def test_low_rank_half_normal(backend, gradient_backend):
     with pm.Model() as model:
-        pm.HalfNormal("a", shape=13)
+        pm.HalfNormal("a", shape=(13, 3))
+        pm.HalfNormal("b", shape=())
+        pm.HalfNormal("c", shape=(5,))
 
     compiled = nutpie.compile_pymc_model(
         model, backend=backend, gradient_backend=gradient_backend