fix: allow variables with zero shapes

Author: aseyboldt

src/pymc.rs

@@ -15,6 +15,7 @@ use pyo3::{
     Bound, Py, PyAny, PyObject, PyResult, Python,
 };
 
+use rand_distr::num_traits::CheckedEuclid;
 use thiserror::Error;
 
 type UserData = *const std::ffi::c_void;
@@ -128,8 +129,8 @@ impl CpuLogpFunc for &LogpFunc {
         let retcode = unsafe {
             (self.func)(
                 self.dim,
-                &position[0] as *const f64,
-                &mut gradient[0] as *mut f64,
+                position.as_ptr(),
+                gradient.as_mut_ptr(),
                 logp_ptr,
                 self.user_data_ptr,
             )
@@ -148,6 +149,7 @@ pub(crate) struct PyMcTrace<'model> {
     var_sizes: Vec<usize>,
     var_names: Vec<String>,
     expand: &'model ExpandFunc,
+    count: usize,
 }
 
 impl<'model> DrawStorage for PyMcTrace<'model> {
@@ -165,14 +167,20 @@ impl<'model> DrawStorage for PyMcTrace<'model> {
             data.extend_from_slice(vals);
             start = end;
         }
+        self.count += 1;
+
         Ok(())
     }
 
     fn finalize(self) -> Result<Arc<dyn Array>> {
         let (fields, arrays): (Vec<_>, _) = izip!(self.data, self.var_names, self.var_sizes)
             .map(|(data, name, size)| {
-                assert!(data.len() % size == 0);
-                let num_arrays = data.len() / size;
+                let (num_arrays, rem) = data
+                    .len()
+                    .checked_div_rem_euclid(&size)
+                    .unwrap_or((self.count, 0));
+                assert!(rem == 0);
+                assert!(num_arrays == self.count);
                 let data = Float64Array::from(data);
                 let item_field = Arc::new(Field::new("item", DataType::Float64, false));
                 let offsets = OffsetBuffer::from_lengths((0..num_arrays).map(|_| size));
@@ -206,6 +214,7 @@ impl<'model> PyMcTrace<'model> {
             var_sizes: model.var_sizes.clone(),
             var_names: model.var_names.clone(),
             expand: &model.expand,
+            count: 0,
         }
     }
 

src/stan.rs

@@ -86,6 +86,9 @@ pub struct StanModel {
 
 /// Return meta information about the constrained parameters of the model
 fn params(var_string: &str) -> anyhow::Result<Vec<Parameter>> {
+    if var_string.is_empty() {
+        return Ok(vec![]);
+    }
     // Parse each variable string into (name, is_complex, indices)
     let parsed_variables: anyhow::Result<Vec<(String, bool, Vec<usize>)>> = var_string
         .split(',')
@@ -540,6 +543,7 @@ pub struct StanTrace<'model> {
     trace: Vec<Vec<f64>>,
     expanded_buffer: Box<[f64]>,
     rng: bridgestan::Rng<&'model bridgestan::StanLibrary>,
+    count: usize,
 }
 
 impl<'model> Clone for StanTrace<'model> {
@@ -559,6 +563,7 @@ impl<'model> Clone for StanTrace<'model> {
             trace: self.trace.clone(),
             expanded_buffer: self.expanded_buffer.clone(),
             rng,
+            count: self.count,
         }
     }
 }
@@ -591,6 +596,7 @@ impl<'model> DrawStorage for StanTrace<'model> {
             // We need to transpose
             fortran_to_c_order(slice, &var.shape, trace);
         }
+        self.count += 1;
         Ok(())
     }
 
@@ -613,7 +619,7 @@ impl<'model> DrawStorage for StanTrace<'model> {
             .unzip();
 
         Ok(Arc::new(
-            StructArray::try_new(fields.into(), arrays, None)
+            StructArray::try_new_with_length(fields.into(), arrays, None, self.count)
                 .context("Could not create arrow StructArray")?,
         ))
     }
@@ -649,6 +655,7 @@ impl Model for StanModel {
             trace,
             rng,
             expanded_buffer: buffer.into(),
+            count: 0,
         })
     }
 
@@ -734,6 +741,10 @@ mod tests {
 
     #[test]
     fn parse_vars() {
+        let vars = "";
+        let parsed = super::params(vars).unwrap();
+        assert!(parsed.len() == 0);
+
         let vars = "x.1.1,x.2.1,x.3.1,x.1.2,x.2.2,x.3.2";
         let parsed = super::params(vars).unwrap();
         assert!(parsed.len() == 1);

tests/test_pymc.py

@@ -31,6 +31,23 @@ def test_pymc_model(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
+@parameterize_backends
+def test_zero_size(backend, gradient_backend):
+    import pytensor.tensor as pt
+
+    with pm.Model() as model:
+        a = pm.Normal("a", shape=(0, 0, 10))
+        pm.Deterministic("b", pt.exp(a))
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    trace = nutpie.sample(compiled, chains=1, draws=17, tune=100)
+    assert trace.posterior.a.shape == (1, 17, 0, 0, 10)
+    assert trace.posterior.b.shape == (1, 17, 0, 0, 10)
+
+
 @pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_float32(backend, gradient_backend):

tests/test_stan.py

@@ -27,6 +27,25 @@ def test_stan_model():
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.stan
+def test_empty():
+    model = """
+    data {}
+    parameters {
+        array[0] real a;
+    }
+    model {
+        a ~ normal(0, 1);
+    }
+    """
+
+    compiled_model = nutpie.compile_stan_model(code=model)
+    trace = nutpie.sample(compiled_model)  # noqa: F841
+    # TODO: Variable `a` is missing because of this bridgestan issue:
+    # https://github.com/roualdes/bridgestan/issues/278
+    # assert trace.posterior.a.shape == (0, 1000)
+
+
 @pytest.mark.stan
 def test_seed():
     model = """