prepare V0.16.1 and fix sampler.abort

CHANGELOG.md

@@ -2,7 +2,26 @@
 
 All notable changes to this project will be documented in this file.
 
-## [0.16.0] - 2025-10-08
+## [0.16.1] - 2025-10-10
+
+### Bug Fixes
+
+- Sampler.abort was broken in last release (Adrian Seyboldt)
+
+- Don't store unconstrained_draw if not requested (Adrian Seyboldt)
+
+
+### Features
+
+- Expose step_size_jitter option (Adrian Seyboldt)
+
+
+### Miscellaneous Tasks
+
+- Cargo update (Adrian Seyboldt)
+
+
+## [0.16.0] - 2025-10-09
 
 ### Bug Fixes
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nutpie"
-version = "0.16.0"
+version = "0.16.1"
 authors = [
   "Adrian Seyboldt <[email protected]>",
   "PyMC Developers <[email protected]>",

python/nutpie/sample.py

@@ -538,7 +538,7 @@ def _extract(self, results):
                 skip_vars = []
                 skips = {
                     "store_gradient": ["gradient"],
-                    "store_unconstrained": ["unconstrained"],
+                    "store_unconstrained": ["unconstrained_draw"],
                     "store_mass_matrix": [
                         "mass_matrix_inv",
                         "mass_matrix_eigvals",
@@ -590,7 +590,7 @@ def is_finished(self):
     def abort(self):
         """Abort sampling and return the trace produced so far."""
         self._sampler.abort()
-        results = self._sampler.extract_results()
+        results = self._sampler.take_results()
         return self._extract(results)
 
     def cancel(self):

src/wrapper.rs

@@ -784,6 +784,35 @@ impl PyNutsSettings {
         };
         Ok(())
     }
+
+    #[getter(step_size_jitter)]
+    fn step_size_jitter(&self) -> Option<f64> {
+        match &self.inner {
+            Settings::LowRank(inner) => inner.adapt_options.step_size_settings.jitter,
+            Settings::Diag(inner) => inner.adapt_options.step_size_settings.jitter,
+            Settings::Transforming(inner) => inner.adapt_options.step_size_settings.jitter,
+        }
+    }
+
+    #[setter(step_size_jitter)]
+    fn set_step_size_jitter(&mut self, mut val: Option<f64>) -> PyResult<()> {
+        if let Some(val) = val {
+            if val < 0.0 {
+                return Err(PyValueError::new_err("step_size_jitter must be positive"));
+            }
+        }
+        if let Some(jitter) = val {
+            if jitter == 0.0 {
+                val = None;
+            }
+        }
+        match &mut self.inner {
+            Settings::LowRank(inner) => inner.adapt_options.step_size_settings.jitter = val,
+            Settings::Diag(inner) => inner.adapt_options.step_size_settings.jitter = val,
+            Settings::Transforming(inner) => inner.adapt_options.step_size_settings.jitter = val,
+        }
+        Ok(())
+    }
 }
 
 pub(crate) enum SamplerState {

tests/test_pymc.py

@@ -192,7 +192,7 @@ def test_wait_timeout(backend, gradient_backend):
 @pytest.mark.timeout(20)
 def test_pause(backend, gradient_backend):
     with pm.Model() as model:
-        pm.Normal("a", shape=100_000)
+        pm.Normal("a", shape=10_000)
     compiled = nutpie.compile_pymc_model(
         model, backend=backend, gradient_backend=gradient_backend
     )
@@ -204,6 +204,23 @@ def test_pause(backend, gradient_backend):
     assert start - time.time() < 5
 
 
+@pytest.mark.pymc
+@parameterize_backends
+@pytest.mark.timeout(20)
+def test_abort(backend, gradient_backend):
+    with pm.Model() as model:
+        pm.Normal("a", shape=10_000)
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    start = time.time()
+    sampler = nutpie.sample(compiled, chains=1, blocking=False)
+    sampler.pause()
+    sampler.resume()
+    sampler.abort()
+    assert start - time.time() < 5
+
+
 @pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_with_coordinate(backend, gradient_backend):
@@ -421,7 +438,6 @@ def test_missing(backend, gradient_backend):
         model, backend=backend, gradient_backend=gradient_backend
     )
     tr = nutpie.sample(compiled, chains=1, seed=1)
-    print(tr.posterior)
     assert hasattr(tr.posterior, "y_unobserved")