Repository: pymc-devs/nutpie (language: python)

Issue:
After chains in warmup go to zero, show how many chains are still sampling instead



Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index 618be46..cab0858 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4,31 +4,26 @@ version = 3
 
 [[package]]
 name = "ahash"
-version = "0.8.3"
+version = "0.8.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2c99f64d1e06488f620f932677e24bc6e2897582980441ae90a671415bd7ec2f"
+checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
 dependencies = [
  "cfg-if",
  "getrandom",
  "once_cell",
  "version_check",
+ "zerocopy",
 ]
 
 [[package]]
 name = "aho-corasick"
-version = "1.0.4"
+version = "1.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6748e8def348ed4d14996fa801f4122cd763fff530258cdc03f64b25f89d3a5a"
+checksum = "b2969dcb958b36655471fc61f7e416fa76033bdd4bfed0678d8fee1e2d07a1f0"
 dependencies = [
  "memchr",
 ]
 
-[[package]]
-name = "android-tzdata"
-version = "0.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
-
 [[package]]
 name = "anes"
 version = "0.1.6"
@@ -37,15 +32,15 @@ checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
 
 [[package]]
 name = "anstyle"
-version = "1.0.2"
+version = "1.0.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "15c4c2c83f81532e5845a733998b6971faca23490340a418e9b72a3ec9de12ea"
+checksum = "8901269c6307e8d93993578286ac0edf7f195079ffff5ebdeea6a59ffb7e36bc"
 
 [[package]]
 name = "anyhow"
-version = "1.0.75"
+version = "1.0.81"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a4668cab20f66d8d020e1fbc0ebe47217433c1b6c8f2040faf858554e394ace6"
+checksum = "0952808a6c2afd1aa8947271f3a60f1a6763c7b912d210184c5149b5cf147247"
 
 [[package]]
 name = "arrow2"
@@ -75,24 +70,24 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
 [[package]]
 name = "bindgen"
-version = "0.66.1"
+version = "0.69.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f2b84e06fc203107bfbad243f4aba2af864eb7db3b1cf46ea0a023b0b433d2a7"
+checksum = "a00dc851838a2120612785d195287475a3ac45514741da670b735818822129a0"
 dependencies = [
- "bitflags 2.4.0",
+ "bitflags 2.4.2",
  "cexpr",
  "clang-sys",
+ "itertools 0.12.1",
  "lazy_static",
  "lazycell",
  "log",
- "peeking_take_while",
  "prettyplease",
  "proc-macro2",
  "quote",
  "regex",
  "rustc-hash",
  "shlex",
- "syn 2.0.29",
+ "syn 2.0.52",
  "which",
 ]
 
@@ -104,45 +99,45 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
 name = "bitflags"
-version = "2.4.0"
+version = "2.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b4682ae6287fcf752ecaabbfcc7b6f9b72aa33933dc23a554d853aea8eea8635"
+checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf"
 
 [[package]]
 name = "bridgestan"
-version = "2.1.2"
+version = "2.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1077a38022365223d0543a3d78666b8680d754542f340a26ab6c79952fa065c"
+checksum = "37a53cdfc0ce7b8511da4a0df5d6aa6fb339c8e42d4d07a2925ed3a60685a50e"
 dependencies = [
  "bindgen",
- "libloading 0.8.0",
+ "libloading",
  "thiserror",
 ]
 
 [[package]]
 name = "bumpalo"
-version = "3.13.0"
+version = "3.15.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1"
+checksum = "7ff69b9dd49fd426c69a0db9fc04dd934cdb6645ff000864d98f7e2af8830eaa"
 
 [[package]]
 name = "bytemuck"
-version = "1.13.1"
+version = "1.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "17febce684fd15d89027105661fec94afb475cb995fbc59d2865198446ba2eea"
+checksum = "5d6d68c57235a3a081186990eca2867354726650f42f7516ca50c28d6281fd15"
 dependencies = [
  "bytemuck_derive",
 ]
 
 [[package]]
 name = "bytemuck_derive"
-version = "1.4.1"
+version = "1.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fdde5c9cd29ebd706ce1b35600920a33550e402fc998a2e53ad3b42c3c47a192"
+checksum = "4da9a32f3fed317401fa3c862968128267c3106685286e15d5aaa3d7389c2f60"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.29",
+ "syn 2.0.52",
 ]
 
 [[package]]
@@ -151,15 +146,6 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 
-[[package]]
-name = "cc"
-version = "1.0.83"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0"
-dependencies = [
- "libc",
-]
-
 [[package]]
 name = "cexpr"
 version = "0.6.0"
@@ -177,19 +163,18 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "chrono"
-version = "0.4.26"
+version = "0.4.35"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec837a71355b28f6556dbd569b37b3f363091c0bd4b2e735674521b4c5fd9bc5"
+checksum = "8eaf5903dcbc0a39312feb77df2ff4c76387d591b9fc7b04a238dcf8bb62639a"
 dependencies = [
- "android-tzdata",
  "num-traits",
 ]
 
 [[package]]
 name = "ciborium"
-version = "0.2.1"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "effd91f6c78e5a4ace8a5d3c0b6bfaec9e2baaef55f3efc00e45fb2e477ee926"
+checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
 dependencies = [
  "ciborium-io",
  "ciborium-ll",
@@ -198,15 +183,15 @@ dependencies = [
 
 [[package]]
 name = "ciborium-io"
-version = "0.2.1"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cdf919175532b369853f5d5e20b26b43112613fd6fe7aee757e35f7a44642656"
+checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
 
 [[package]]
 name = "ciborium-ll"
-version = "0.2.1"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "defaa24ecc093c77630e6c15e17c51f5e187bf35ee514f4e2d67baaa96dae22b"
+checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
 dependencies = [
  "ciborium-io",
  "half",
@@ -214,29 +199,29 @@ dependencies = [
 
 [[package]]
 name = "clang-sys"
-version = "1.6.1"
+version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c688fc74432808e3eb684cae8830a86be1d66a2bd58e1f248ed0960a590baf6f"
+checksum = "67523a3b4be3ce1989d607a828d036249522dd9c1c8de7f4dd2dae43a37369d1"
 dependencies = [
  "glob",
  "libc",
- "libloading 0.7.4",
+ "libloading",
 ]
 
 [[package]]
 name = "clap"
-version = "4.4.0"
+version = "4.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d5f1946157a96594eb2d2c10eb7ad9a2b27518cb3000209dec700c35df9197d"
+checksum = "949626d00e063efc93b6dca932419ceb5432f99769911c0b995f7e884c778813"
 dependencies = [
  "clap_builder",
 ]
 
 [[package]]
 name = "clap_builder"
-version = "4.4.0"
+version = "4.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78116e32a042dd73c2901f0dc30790d20ff3447f3e3472fad359e8c3d282bcd6"
+checksum = "ae129e2e766ae0ec03484e609954119f123cc1fe650337e155d03b022f24f7b4"
 dependencies = [
  "anstyle",
  "clap_lex",
@@ -244,9 +229,9 @@ dependencies = [
 
 [[package]]
 name = "clap_lex"
-version = "0.5.1"
+version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd7cc57abe963c6d3b9d8be5b06ba7c8957a930305ca90304f24ef040aa6f961"
+checksum = "98cc8fbded0c607b7ba9dd60cd98df59af97e84d24e49c8557331cfc26d301ce"
 
 [[package]]
 name = "criterion"
@@ -286,11 +271,10 @@ dependencies = [
 
 [[package]]
 name = "crossbeam"
-version = "0.8.2"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2801af0d36612ae591caa9568261fddce32ce6e08a7275ea334a06a4ad021a2c"
+checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8"
 dependencies = [
- "cfg-if",
  "crossbeam-channel",
  "crossbeam-deque",
  "crossbeam-epoch",
@@ -300,95 +284,80 @@ dependencies = [
 
 [[package]]
 name = "crossbeam-channel"
-version = "0.5.8"
+version = "0.5.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a33c2bf77f2df06183c3aa30d1e96c0695a313d4f9c453cc3762a6db39f99200"
+checksum = "ab3db02a9c5b5121e1e42fbdb1aeb65f5e02624cc58c43f2884c6ccac0b82f95"
 dependencies = [
- "cfg-if",
  "crossbeam-utils",
 ]
 
 [[package]]
 name = "crossbeam-deque"
-version = "0.8.3"
+version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef"
+checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d"
 dependencies = [
- "cfg-if",
  "crossbeam-epoch",
  "crossbeam-utils",
 ]
 
 [[package]]
 name = "crossbeam-epoch"
-version = "0.9.15"
+version = "0.9.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7"
+checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
 dependencies = [
- "autocfg",
- "cfg-if",
  "crossbeam-utils",
- "memoffset",
- "scopeguard",
 ]
 
 [[package]]
 name = "crossbeam-queue"
-version = "0.3.8"
+version = "0.3.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d1cfb3ea8a53f37c40dea2c7bedcbd88bdfae54f5e2175d6ecaff1c988353add"
+checksum = "df0346b5d5e76ac2fe4e327c5fd1118d6be7c51dfb18f9b7922923f287471e35"
 dependencies = [
- "cfg-if",
  "crossbeam-utils",
 ]
 
 [[package]]
 name = "crossbeam-utils"
-version = "0.8.16"
+version = "0.8.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294"
-dependencies = [
- "cfg-if",
-]
+checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345"
+
+[[package]]
+name = "crunchy"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
 
 [[package]]
 name = "dyn-clone"
-version = "1.0.13"
+version = "1.0.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bbfc4744c1b8f2a09adc0e55242f60b1af195d88596bd8700be74418c056c555"
+checksum = "0d6ef0072f8a535281e4876be788938b528e9a1d43900b82c2569af7da799125"
 
 [[package]]
 name = "either"
-version = "1.9.0"
+version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
+checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a"
 
 [[package]]
 name = "errno"
-version = "0.3.2"
+version = "0.3.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6b30f669a7961ef1631673d2766cc92f52d64f7ef354d4fe0ddfd30ed52f0f4f"
+checksum = "a258e46cdc063eb8519c00b9fc845fc47bcfca4130e2f08e88665ceda8474245"
 dependencies = [
- "errno-dragonfly",
  "libc",
  "windows-sys",
 ]
 
-[[package]]
-name = "errno-dragonfly"
-version = "0.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf"
-dependencies = [
- "cc",
- "libc",
-]
-
 [[package]]
 name = "ethnum"
-version = "1.3.2"
+version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0198b9d0078e0f30dedc7acbb21c974e838fc8fae3ee170128658a98cb2c1c04"
+checksum = "b90ca2580b73ab6a1f724b76ca11ab632df820fd6040c336200d2c1df7b3c82c"
 
 [[package]]
 name = "foreign_vec"
@@ -398,9 +367,9 @@ checksum = "ee1b05cbd864bcaecbd3455d6d967862d446e4ebfc3c2e5e5b9841e53cba6673"
 
 [[package]]
 name = "getrandom"
-version = "0.2.10"
+version = "0.2.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427"
+checksum = "190092ea657667030ac6a35e305e62fc4dd69fd98ac98631e5d3a2b1575a12b5"
 dependencies = [
  "cfg-if",
  "js-sys",
@@ -417,9 +386,13 @@ checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
 
 [[package]]
 name = "half"
-version = "1.8.2"
+version = "2.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
+checksum = "b5eceaaeec696539ddaf7b333340f1af35a5aa87ae3e4f3ead0532f72affab2e"
+dependencies = [
+ "cfg-if",
+ "crunchy",
+]
 
 [[package]]
 name = "hash_hasher"
@@ -427,26 +400,41 @@ version = "2.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "74721d007512d0cb3338cd20f0654ac913920061a4c4d0d8708edb3f2a698c0c"
 
+[[package]]
+name = "heck"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
+
 [[package]]
 name = "hermit-abi"
-version = "0.3.2"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024"
+
+[[package]]
+name = "home"
+version = "0.5.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b"
+checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5"
+dependencies = [
+ "windows-sys",
+]
 
 [[package]]
 name = "indoc"
-version = "1.0.9"
+version = "2.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bfa799dd5ed20a7e349f3b4639aa80d74549c81716d9ec4f994c9b5815598306"
+checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8"
 
 [[package]]
 name = "is-terminal"
-version = "0.4.9"
+version = "0.4.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b"
+checksum = "f23ff5ef2b80d608d61efee834934d862cd92461afc0560dedf493e4c033738b"
 dependencies = [
  "hermit-abi",
- "rustix",
+ "libc",
  "windows-sys",
 ]
 
@@ -468,17 +456,26 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
-version = "1.0.9"
+version = "1.0.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38"
+checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c"
 
 [[package]]
 name = "js-sys"
-version = "0.3.64"
+version = "0.3.69"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c5f195fe497f702db0f318b07fdd68edb16955aed830df8363d837542f8f935a"
+checksum = "29c15563dc2726973df627357ce0c9ddddbea194836909d655df6a75d2cf296d"
 dependencies = [
  "wasm-bindgen",
 ]
@@ -497,35 +494,25 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55"
 
 [[package]]
 name = "libc"
-version = "0.2.147"
+version = "0.2.153"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b4668fb0ea861c1df094127ac5f1da3409a82116a4ba74fca2e58ef927159bb3"
+checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd"
 
 [[package]]
 name = "libloading"
-version = "0.7.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f"
-dependencies = [
- "cfg-if",
- "winapi",
-]
-
-[[package]]
-name = "libloading"
-version = "0.8.0"
+version = "0.8.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d580318f95776505201b28cf98eb1fa5e4be3b689633ba6a3e6cd880ff22d8cb"
+checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19"
 dependencies = [
  "cfg-if",
- "windows-sys",
+ "windows-targets 0.52.4",
 ]
 
 [[package]]
 name = "libm"
-version = "0.2.7"
+version = "0.2.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4"
+checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
 
 [[package]]
 name = "linux-raw-sys"
@@ -535,9 +522,9 @@ checksum = "01cda141df6706de531b6c46c3a33ecca755538219bd484262fa09410c13539c"
 
 [[package]]
 name = "lock_api"
-version = "0.4.10"
+version = "0.4.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c1cc9717a20b1bb222f333e6a92fd32f7d8a18ddc5a3191a11af45dcbf4dcd16"
+checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45"
 dependencies = [
  "autocfg",
  "scopeguard",
@@ -545,15 +532,15 @@ dependencies = [
 
 [[package]]
 name = "log"
-version = "0.4.20"
+version = "0.4.21"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
+checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c"
 
 [[package]]
 name = "matrixmultiply"
-version = "0.3.7"
+version = "0.3.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "090126dc04f95dc0d1c1c91f61bdd474b3930ca064c1edc8a849da2c6cbe1e77"
+checksum = "7574c1cf36da4798ab73da5b215bbf444f50718207754cb522201d78d1cd0ff2"
 dependencies = [
  "autocfg",
  "rawpointer",
@@ -561,9 +548,9 @@ dependencies = [
 
 [[package]]
 name = "memchr"
-version = "2.5.0"
+version = "2.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
+checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149"
 
 [[package]]
 name = "memoffset"
@@ -582,9 +569,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
 
 [[package]]
 name = "multiversion"
-version = "0.7.3"
+version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b2c7b9d7fe61760ce5ea19532ead98541f6b4c495d87247aff9826445cf6872a"
+checksum = "c4851161a11d3ad0bf9402d90ffc3967bf231768bfd7aeb61755ad06dbf1a142"
 dependencies = [
  "multiversion-macros",
  "target-features",
@@ -592,9 +579,9 @@ dependencies = [
 
 [[package]]
 name = "multiversion-macros"
-version = "0.7.3"
+version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "26a83d8500ed06d68877e9de1dde76c1dbb83885dcdbda4ef44ccbc3fbda2ac8"
+checksum = "79a74ddee9e0c27d2578323c13905793e91622148f138ba29738f9dddb835e90"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -627,48 +614,37 @@ dependencies = [
 
 [[package]]
 name = "num-complex"
-version = "0.4.4"
+version = "0.4.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214"
+checksum = "23c6602fda94a57c990fe0df199a035d83576b496aa29f4e634a8ac6004e68a6"
 dependencies = [
  "num-traits",
 ]
 
 [[package]]
 name = "num-integer"
-version = "0.1.45"
+version = "0.1.46"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
 dependencies = [
- "autocfg",
  "num-traits",
 ]
 
 [[package]]
 name = "num-traits"
-version = "0.2.16"
+version = "0.2.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f30b0abd723be7e2ffca1272140fac1a2f084c77ec3e123c192b66af1ee9e6c2"
+checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a"
 dependencies = [
  "autocfg",
  "libm",
 ]
 
-[[package]]
-name = "num_cpus"
-version = "1.16.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43"
-dependencies = [
- "hermit-abi",
- "libc",
-]
-
 [[package]]
 name = "numpy"
-version = "0.19.0"
+version = "0.20.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "437213adf41bbccf4aeae535fbfcdad0f6fed241e1ae182ebe97fa1f3ce19389"
+checksum = "bef41cbb417ea83b30525259e30ccef6af39b31c240bda578889494c5392d331"
 dependencies = [
  "libc",
  "ndarray",
@@ -687,7 +663,7 @@ dependencies = [
  "arrow2",
  "bridgestan",
  "criterion",
- "itertools 0.11.0",
+ "itertools 0.12.1",
  "ndarray",
  "numpy",
  "nuts-rs",
@@ -720,9 +696,9 @@ dependencies = [
 
 [[package]]
 name = "once_cell"
-version = "1.18.0"
+version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
+checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
 [[package]]
 name = "oorandom"
@@ -742,23 +718,17 @@ dependencies = [
 
 [[package]]
 name = "parking_lot_core"
-version = "0.9.8"
+version = "0.9.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93f00c865fe7cabf650081affecd3871070f26767e7b2070a3ffae14c654b447"
+checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e"
 dependencies = [
  "cfg-if",
  "libc",
  "redox_syscall",
  "smallvec",
- "windows-targets",
+ "windows-targets 0.48.5",
 ]
 
-[[package]]
-name = "peeking_take_while"
-version = "0.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099"
-
 [[package]]
 name = "plotters"
 version = "0.3.5"
@@ -787,6 +757,12 @@ dependencies = [
  "plotters-backend",
 ]
 
+[[package]]
+name = "portable-atomic"
+version = "1.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0"
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.17"
@@ -795,28 +771,28 @@ checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
 
 [[package]]
 name = "prettyplease"
-version = "0.2.12"
+version = "0.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6c64d9ba0963cdcea2e1b2230fbae2bab30eb25a174be395c41e764bfb65dd62"
+checksum = "a41cf62165e97c7f814d2221421dbb9afcbcdb0a88068e5ea206e19951c2cbb5"
 dependencies = [
  "proc-macro2",
- "syn 2.0.29",
+ "syn 2.0.52",
 ]
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.66"
+version = "1.0.79"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "18fb31db3f9bddb2ea821cde30a9f70117e3f119938b5ee630b7403aa6e2ead9"
+checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e"
 dependencies = [
  "unicode-ident",
 ]
 
 [[package]]
 name = "pyo3"
-version = "0.19.2"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e681a6cfdc4adcc93b4d3cf993749a4552018ee0a9b65fc0ccfad74352c72a38"
+checksum = "53bdbb96d49157e65d45cc287af5f32ffadd5f4761438b527b055fb0d4bb8233"
 dependencies = [
  "anyhow",
  "cfg-if",
@@ -824,6 +800,7 @@ dependencies = [
  "libc",
  "memoffset",
  "parking_lot",
+ "portable-atomic",
  "pyo3-build-config",
  "pyo3-ffi",
  "pyo3-macros",
@@ -832,9 +809,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.19.2"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "076c73d0bc438f7a4ef6fdd0c3bb4732149136abd952b110ac93e4edb13a6ba5"
+checksum = "deaa5745de3f5231ce10517a1f5dd97d53e5a2fd77aa6b5842292085831d48d7"
 dependencies = [
  "once_cell",
  "target-lexicon",
@@ -842,9 +819,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.19.2"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e53cee42e77ebe256066ba8aa77eff722b3bb91f3419177cf4cd0f304d3284d9"
+checksum = "62b42531d03e08d4ef1f6e85a2ed422eb678b8cd62b762e53891c05faf0d4afa"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -852,32 +829,34 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.19.2"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dfeb4c99597e136528c6dd7d5e3de5434d1ceaf487436a3f03b2d56b6fc9efd1"
+checksum = "7305c720fa01b8055ec95e484a6eca7a83c841267f0dd5280f0c8b8551d2c158"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
  "quote",
- "syn 1.0.109",
+ "syn 2.0.52",
 ]
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.19.2"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "947dc12175c254889edc0c02e399476c2f652b4b9ebd123aa655c224de259536"
+checksum = "7c7e9b68bb9c3149c5b0cade5d07f953d6d125eb4337723c4ccdb665f1f96185"
 dependencies = [
+ "heck",
  "proc-macro2",
+ "pyo3-build-config",
  "quote",
- "syn 1.0.109",
+ "syn 2.0.52",
 ]
 
 [[package]]
 name = "quote"
-version = "1.0.33"
+version = "1.0.35"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae"
+checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef"
 dependencies = [
  "proc-macro2",
 ]
@@ -930,9 +909,9 @@ checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
 
 [[package]]
 name = "rayon"
-version = "1.7.0"
+version = "1.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d2df5196e37bcc87abebc0053e20787d73847bb33134a69841207dd0a47f03b"
+checksum = "e4963ed1bc86e4f3ee217022bd855b297cef07fb9eac5dfa1f788b220b49b3bd"
 dependencies = [
  "either",
  "rayon-core",
@@ -940,30 +919,28 @@ dependencies = [
 
 [[package]]
 name = "rayon-core"
-version = "1.11.0"
+version = "1.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b8f95bd6966f5c87776639160a66bd8ab9895d9d4ab01ddba9fc60661aebe8d"
+checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2"
 dependencies = [
- "crossbeam-channel",
  "crossbeam-deque",
  "crossbeam-utils",
- "num_cpus",
 ]
 
 [[package]]
 name = "redox_syscall"
-version = "0.3.5"
+version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "567664f262709473930a4bf9e51bf2ebf3348f2e748ccc50dea20646858f8f29"
+checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa"
 dependencies = [
  "bitflags 1.3.2",
 ]
 
 [[package]]
 name = "regex"
-version = "1.9.3"
+version = "1.10.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "81bc1d4caf89fac26a70747fe603c130093b53c773888797a6329091246d651a"
+checksum = "b62dbe01f0b06f9d8dc7d49e05a0785f153b00b2c227856282f671e0318c9b15"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -973,9 +950,9 @@ dependencies = [
 
 [[package]]
 name = "regex-automata"
-version = "0.3.6"
+version = "0.4.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fed1ceff11a1dddaee50c9dc8e4938bd106e9d89ae372f192311e7da498e3b69"
+checksum = "86b83b8b9847f9bf95ef68afb0b8e6cdb80f498442f5179a29fad448fcc1eaea"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -984,9 +961,9 @@ dependencies = [
 
 [[package]]
 name = "regex-syntax"
-version = "0.7.4"
+version = "0.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e5ea92a5b6195c6ef2a0295ea818b312502c6fc94dde986c5553242e18fd4ce2"
+checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
 
 [[package]]
 name = "rustc-hash"
@@ -1005,11 +982,11 @@ dependencies = [
 
 [[package]]
 name = "rustix"
-version = "0.38.13"
+version = "0.38.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7db8590df6dfcd144d22afd1b83b36c21a18d7cbc1dc4bb5295a8712e9eb662"
+checksum = "6ea3e1a662af26cd7a3ba09c0297a31af215563ecf42817c98df621387f4e949"
 dependencies = [
- "bitflags 2.4.0",
+ "bitflags 2.4.2",
  "errno",
  "libc",
  "linux-raw-sys",
@@ -1018,9 +995,9 @@ dependencies = [
 
 [[package]]
 name = "ryu"
-version = "1.0.15"
+version = "1.0.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741"
+checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1"
 
 [[package]]
 name = "same-file"
@@ -1039,35 +1016,35 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
 
 [[package]]
 name = "semver"
-version = "1.0.18"
+version = "1.0.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0293b4b29daaf487284529cc2f5675b8e57c61f70167ba415a463651fd6a918"
+checksum = "92d43fe69e652f3df9bdc2b85b2854a0825b86e4fb76bc44d945137d053639ca"
 
 [[package]]
 name = "serde"
-version = "1.0.186"
+version = "1.0.197"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f5db24220c009de9bd45e69fb2938f4b6d2df856aa9304ce377b3180f83b7c1"
+checksum = "3fb1c873e1b9b056a4dc4c0c198b24c3ffa059243875552b2bd0933b1aee4ce2"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.186"
+version = "1.0.197"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5ad697f7e0b65af4983a4ce8f56ed5b357e8d3c36651bf6a7e13639c17b8e670"
+checksum = "7eb0b34b42edc17f6b7cac84a52a1c5f0e1bb2227e997ca9011ea3dd34e8610b"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.29",
+ "syn 2.0.52",
 ]
 
 [[package]]
 name = "serde_json"
-version = "1.0.105"
+version = "1.0.114"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "693151e1ac27563d6dbcec9dee9fbd5da8539b20fa14ad3752b2e6d363ace360"
+checksum = "c5f09b1bd632ef549eaa9f60a1f8de742bdbc698e6cee2095fc84dde5f549ae0"
 dependencies = [
  "itoa",
  "ryu",
@@ -1088,9 +1065,9 @@ checksum = "f27f6278552951f1f2b8cf9da965d10969b2efdea95a6ec47987ab46edfe263a"
 
 [[package]]
 name = "smallvec"
-version = "1.11.0"
+version = "1.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9"
+checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7"
 
 [[package]]
 name = "syn"
@@ -1105,9 +1082,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.29"
+version = "2.0.52"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c324c494eba9d92503e6f1ef2e6df781e78f6a7705a0202d9801b198807d518a"
+checksum = "b699d15b36d1f02c3e7c69f8ffef53de37aefae075d8488d4ba1a7788d574a07"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -1116,34 +1093,34 @@ dependencies = [
 
 [[package]]
 name = "target-features"
-version = "0.1.4"
+version = "0.1.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "06f6b473c37f9add4cf1df5b4d66a8ef58ab6c895f1a3b3f949cf3e21230140e"
+checksum = "c1bbb9f3c5c463a01705937a24fdabc5047929ac764b2d5b9cf681c1f5041ed5"
 
 [[package]]
 name = "target-lexicon"
-version = "0.12.11"
+version = "0.12.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9d0e916b1148c8e263850e1ebcbd046f333e0683c724876bb0da63ea4373dc8a"
+checksum = "e1fc403891a21bcfb7c37834ba66a547a8f402146eba7265b5a6d88059c9ff2f"
 
 [[package]]
 name = "thiserror"
-version = "1.0.47"
+version = "1.0.58"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97a802ec30afc17eee47b2855fc72e0c4cd62be9b4efe6591edde0ec5bd68d8f"
+checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297"
 dependencies = [
  "thiserror-impl",
 ]
 
 [[package]]
 name = "thiserror-impl"
-version = "1.0.47"
+version = "1.0.58"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6bb623b56e39ab7dcd4b1b98bb6c8f8d907ed255b18de254088016b27a8ee19b"
+checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.29",
+ "syn 2.0.52",
 ]
 
 [[package]]
@@ -1158,15 +1135,15 @@ dependencies = [
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.11"
+version = "1.0.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c"
+checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
 
 [[package]]
 name = "unindent"
-version = "0.1.11"
+version = "0.2.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e1766d682d402817b5ac4490b3c3002d91dfa0d22812f341609f97b08757359c"
+checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce"
 
 [[package]]
 name = "version_check"
@@ -1176,9 +1153,9 @@ checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 
 [[package]]
 name = "walkdir"
-version = "2.3.3"
+version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "36df944cda56c7d8d8b7496af378e6b16de9284591917d307c9b4d313c44e698"
+checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
 dependencies = [
  "same-file",
  "winapi-util",
@@ -1192,9 +1169,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.87"
+version = "0.2.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7706a72ab36d8cb1f80ffbf0e071533974a60d0a308d01a5d0375bf60499a342"
+checksum = "4be2531df63900aeb2bca0daaaddec08491ee64ceecbee5076636a3b026795a8"
 dependencies = [
  "cfg-if",
  "wasm-bindgen-macro",
@@ -1202,24 +1179,24 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-backend"
-version = "0.2.87"
+version = "0.2.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5ef2b6d3c510e9625e5fe6f509ab07d66a760f0885d858736483c32ed7809abd"
+checksum = "614d787b966d3989fa7bb98a654e369c762374fd3213d212cfc0251257e747da"
 dependencies = [
  "bumpalo",
  "log",
  "once_cell",
  "proc-macro2",
  "quote",
- "syn 2.0.29",
+ "syn 2.0.52",
  "wasm-bindgen-shared",
 ]
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.87"
+version = "0.2.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dee495e55982a3bd48105a7b947fd2a9b4a8ae3010041b9e0faab3f9cd028f1d"
+checksum = "a1f8823de937b71b9460c0c34e25f3da88250760bec0ebac694b49997550d726"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -1227,28 +1204,28 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.87"
+version = "0.2.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b"
+checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.29",
+ "syn 2.0.52",
  "wasm-bindgen-backend",
  "wasm-bindgen-shared",
 ]
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.87"
+version = "0.2.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ca6ad05a4870b2bf5fe995117d3728437bd27d7cd5f06f13c17443ef369775a1"
+checksum = "af190c94f2773fdb3729c55b007a722abb5384da03bc0986df4c289bf5567e96"
 
 [[package]]
 name = "web-sys"
-version = "0.3.64"
+version = "0.3.69"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9b85cbef8c220a6abc02aefd892dfc0fc23afb1c6a426316ec33253a3877249b"
+checksum = "77afa9a11836342370f4817622a2f0f418b134426d91a82dfb48f532d2ec13ef"
 dependencies = [
  "js-sys",
  "wasm-bindgen",
@@ -1256,13 +1233,14 @@ dependencies = [
 
 [[package]]
 name = "which"
-version = "4.4.0"
+version = "4.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2441c784c52b289a054b7201fc93253e288f094e2f4be9058343127c4226a269"
+checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7"
 dependencies = [
  "either",
- "libc",
+ "home",
  "once_cell",
+ "rustix",
 ]
 
 [[package]]
@@ -1283,9 +1261,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
 
 [[package]]
 name = "winapi-util"
-version = "0.1.5"
+version = "0.1.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178"
+checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596"
 dependencies = [
  "winapi",
 ]
@@ -1298,11 +1276,11 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
 [[package]]
 name = "windows-sys"
-version = "0.48.0"
+version = "0.52.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9"
+checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
 dependencies = [
- "windows-targets",
+ "windows-targets 0.52.4",
 ]
 
 [[package]]
@@ -1311,13 +1289,28 @@ version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c"
 dependencies = [
- "windows_aarch64_gnullvm",
- "windows_aarch64_msvc",
- "windows_i686_gnu",
- "windows_i686_msvc",
- "windows_x86_64_gnu",
- "windows_x86_64_gnullvm",
- "windows_x86_64_msvc",
+ "windows_aarch64_gnullvm 0.48.5",
+ "windows_aarch64_msvc 0.48.5",
+ "windows_i686_gnu 0.48.5",
+ "windows_i686_msvc 0.48.5",
+ "windows_x86_64_gnu 0.48.5",
+ "windows_x86_64_gnullvm 0.48.5",
+ "windows_x86_64_msvc 0.48.5",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7dd37b7e5ab9018759f893a1952c9420d060016fc19a472b4bb20d1bdd694d1b"
+dependencies = [
+ "windows_aarch64_gnullvm 0.52.4",
+ "windows_aarch64_msvc 0.52.4",
+ "windows_i686_gnu 0.52.4",
+ "windows_i686_msvc 0.52.4",
+ "windows_x86_64_gnu 0.52.4",
+ "windows_x86_64_gnullvm 0.52.4",
+ "windows_x86_64_msvc 0.52.4",
 ]
 
 [[package]]
@@ -1326,38 +1319,104 @@ version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8"
 
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bcf46cf4c365c6f2d1cc93ce535f2c8b244591df96ceee75d8e83deb70a9cac9"
+
 [[package]]
 name = "windows_aarch64_msvc"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc"
 
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da9f259dd3bcf6990b55bffd094c4f7235817ba4ceebde8e6d11cd0c5633b675"
+
 [[package]]
 name = "windows_i686_gnu"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e"
 
+[[package]]
+name = "windows_i686_gnu"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b474d8268f99e0995f25b9f095bc7434632601028cf86590aea5c8a5cb7801d3"
+
 [[package]]
 name = "windows_i686_msvc"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406"
 
+[[package]]
+name = "windows_i686_msvc"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1515e9a29e5bed743cb4415a9ecf5dfca648ce85ee42e15873c3cd8610ff8e02"
+
 [[package]]
 name = "windows_x86_64_gnu"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e"
 
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5eee091590e89cc02ad514ffe3ead9eb6b660aedca2183455434b93546371a03"
+
 [[package]]
 name = "windows_x86_64_gnullvm"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc"
 
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ca79f2451b49fa9e2af39f0747fe999fcda4f5e241b2898624dca97a1f2177"
+
 [[package]]
 name = "windows_x86_64_msvc"
 version = "0.48.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.52.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32b752e52a2da0ddfbdbcc6fceadfeede4c939ed16d13e648833a61dfb611ed8"
+
+[[package]]
+name = "zerocopy"
+version = "0.7.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74d4d3961e53fa4c9a25a8637fc2bfaf2595b3d3ae34875568a5cf64787716be"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.7.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ce1b18ccd8e73a9321186f97e46f9f04b778851177567b1975109d26a08d2a6"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.52",
+]
+
+[[patch.unused]]
+name = "eerie"
+version = "0.2.1"
diff --git a/Cargo.toml b/Cargo.toml
index 64d792b..a113017 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,21 +22,21 @@ crate-type = ["cdylib"]
 
 [dependencies]
 nuts-rs = "0.7.0"
-numpy = "0.19.0"
+numpy = "0.20.0"
 ndarray = "0.15.6"
 rand = "0.8.5"
 thiserror = "1.0.44"
 rand_chacha = "0.3.1"
-rayon = "1.7.0"
-arrow2 = "0.17.3"
+rayon = "1.9.0"
+arrow2 = "0.17.0"
 anyhow = "1.0.72"
-itertools = "0.11.0"
+itertools = "0.12.0"
 bridgestan = "2.1.2"
 rand_distr = "0.4.3"
 smallvec = "1.11.0"
 
 [dependencies.pyo3]
-version = "0.19.2"
+version = "0.20.0"
 features = ["extension-module", "anyhow"]
 
 [dev-dependencies]
diff --git a/README.md b/README.md
index 7f4ef9d..58e3f06 100644
--- a/README.md
+++ b/README.md
@@ -94,6 +94,26 @@ trace_pymc = nutpie.sample(compiled_model)
 `trace_pymc` now contains an ArviZ `InferenceData` object, including sampling
 statistics and the posterior of the variables defined above.
 
+We can also control the sampler in a non-blocking way:
+
+```python
+# The sampler will now run the the background
+sampler = nutpie.sample(compiled_model, blocking=False)
+
+# Pause and resume the sampling
+sampler.pause()
+sampler.resume()
+
+# Wait for the sampler to finish (up to timeout seconds)
+# sampler.wait(timeout=0.1)
+
+# or we can also abort the sampler (and return the incomplete trace)
+incomplete_trace = sampler.abort()
+
+# or cancel and discard all progress:
+sampler.cancel()
+```
+
 ## Usage with Stan
 
 In order to sample from Stan model, `bridgestan` needs to be installed.
diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py
index e8d3016..cd32b15 100644
--- a/python/nutpie/compile_pymc.py
+++ b/python/nutpie/compile_pymc.py
@@ -86,9 +86,9 @@ def with_data(self, **updates):
             user_data=user_data,
         )
 
-    def _make_sampler(self, settings, init_mean, chains, cores, seed):
+    def _make_sampler(self, settings, init_mean, chains, cores, seed, callback=None):
         model = self._make_model(init_mean)
-        return _lib.PySampler.from_pymc(settings, chains, cores, model, seed)
+        return _lib.PySampler.from_pymc(settings, chains, cores, model, seed, callback)
 
     def _make_model(self, init_mean):
         expand_fn = _lib.ExpandFunc(
diff --git a/python/nutpie/compile_stan.py b/python/nutpie/compile_stan.py
index 2d0aef4..8cf86a3 100644
--- a/python/nutpie/compile_stan.py
+++ b/python/nutpie/compile_stan.py
@@ -80,9 +80,9 @@ def _make_model(self, init_mean):
             return self.with_data().model
         return self.model
 
-    def _make_sampler(self, settings, init_mean, chains, cores, seed):
+    def _make_sampler(self, settings, init_mean, chains, cores, seed, callback=None):
         model = self._make_model(init_mean)
-        return _lib.PySampler.from_stan(settings, chains, cores, model, seed)
+        return _lib.PySampler.from_stan(settings, chains, cores, model, seed, callback)
 
     @property
     def n_dim(self):
diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py
index 3cbcf96..9a97843 100644
--- a/python/nutpie/sample.py
+++ b/python/nutpie/sample.py
@@ -1,11 +1,13 @@
 from dataclasses import dataclass
-from typing import Optional
+from threading import Condition, Event
+from typing import Any, Literal, Optional, overload
 
 import arviz
 import fastprogress
 import numpy as np
 import pandas as pd
 import pyarrow
+from fastprogress.fastprogress import ConsoleProgressBar
 
 from nutpie import _lib
 
@@ -125,6 +127,203 @@ def _trace_to_arviz(traces, n_tune, shapes, **kwargs):
     )
 
 
+class _BackgroundSampler:
+    _sampler: Any
+    _num_divs: int
+    _tune: int
+    _draws: int
+    _chains: int
+    _chains_finished: int
+    _compiled_model: CompiledModel
+    _save_warmup: bool
+    _progress = ConsoleProgressBar
+
+    def __init__(
+        self,
+        compiled_model,
+        settings,
+        init_mean,
+        chains,
+        cores,
+        seed,
+        draws,
+        tune,
+        *,
+        progress_bar=True,
+        save_warmup=True,
+        return_raw_trace=False,
+    ):
+        self._num_divs = 0
+        self._tune = settings.num_tune
+        self._draws = settings.num_draws
+        self._settings = settings
+        self._chains_tuning = chains
+        self._chains_finished = 0
+        self._chains = chains
+        self._compiled_model = compiled_model
+        self._save_warmup = save_warmup
+        self._return_raw_trace = return_raw_trace
+        total_draws = (self._draws + self._tune) * self._chains
+        self._progress = fastprogress.progress_bar(
+            range(total_draws),
+            total=total_draws,
+            display=progress_bar,
+        )
+        # fastprogress seems to reset the progress bar
+        # if we create a new iterator, but we don't want
+        # this for multiple calls to wait.
+        self._bar = iter(self._progress)
+
+        self._exit_event = Event()
+        self._pause_event = Event()
+        self._continue = Condition()
+
+        self._finished_draws = 0
+
+        next(self._bar)
+
+        def progress_callback(info):
+            if info.draw == self._tune - 1:
+                self._chains_tuning -= 1
+            if info.draw == self._tune + self._draws - 1:
+                self._chains_finished += 1
+            if info.is_diverging and info.draw > self._tune:
+                self._num_divs += 1
+            if self._chains_tuning > 0:
+                count = self._chains_tuning
+                divs = self._num_divs
+                self._progress.comment = (
+                    f" Chains in warmup: {count}, Divergences: {divs}"
+                )
+            else:
+                count = self._chains - self._chains_finished
+                divs = self._num_divs
+                self._progress.comment = (
+                    f" Sampling chains: {count}, Divergences: {divs}"
+                )
+            try:
+                next(self._bar)
+            except StopIteration:
+                pass
+            self._finished_draws += 1
+
+        if progress_bar:
+            callback = progress_callback
+        else:
+            callback = None
+
+        self._sampler = compiled_model._make_sampler(
+            settings,
+            init_mean,
+            chains,
+            cores,
+            seed,
+            callback=callback,
+        )
+
+    def wait(self, *, timeout=None):
+        """Wait until sampling is finished and return the trace.
+
+        KeyboardInterrupt will lead to interrupt the waiting.
+
+        This will return after `timeout` seconds even if sampling is
+        not finished at this point.
+
+        This resumes the sampler in case it had been paused.
+        """
+        self._sampler.wait(timeout)
+        self._sampler.finalize()
+        return self._extract()
+
+    def _extract(self):
+        results = self._sampler.extract_results()
+
+        dims = {name: list(dim) for name, dim in self._compiled_model.dims.items()}
+        dims["mass_matrix_inv"] = ["unconstrained_parameter"]
+        dims["gradient"] = ["unconstrained_parameter"]
+        dims["unconstrained_draw"] = ["unconstrained_parameter"]
+        dims["divergence_start"] = ["unconstrained_parameter"]
+        dims["divergence_start_gradient"] = ["unconstrained_parameter"]
+        dims["divergence_end"] = ["unconstrained_parameter"]
+        dims["divergence_momentum"] = ["unconstrained_parameter"]
+
+        if self._return_raw_trace:
+            return results
+        else:
+            return _trace_to_arviz(
+                results,
+                self._tune,
+                self._compiled_model.shapes,
+                dims=dims,
+                coords={
+                    name: pd.Index(vals)
+                    for name, vals in self._compiled_model.coords.items()
+                },
+                save_warmup=self._save_warmup,
+            )
+
+    def pause(self):
+        """Pause the sampler."""
+        self._sampler.pause()
+
+    def resume(self):
+        """Resume a paused sampler."""
+        self._sampler.resume()
+
+    @property
+    def is_finished(self):
+        return self._sampler.is_finished()
+
+    def abort(self):
+        """Abort sampling and return the trace produced so far."""
+        self._sampler.abort()
+        return self._extract()
+
+    def cancel(self):
+        """Abort sampling and discard progress."""
+        self._sampler.abort()
+
+    def __del__(self):
+        if not self._sampler.is_empty():
+            self.cancel()
+
+
+@overload
+def sample(
+    compiled_model: CompiledModel,
+    *,
+    draws: int,
+    tune: int,
+    chains: int,
+    cores: int,
+    seed: Optional[int],
+    save_warmup: bool,
+    progress_bar: bool,
+    init_mean: Optional[np.ndarray],
+    return_raw_trace: bool,
+    blocking: Literal[True],
+    **kwargs,
+) -> arviz.InferenceData: ...
+
+
+@overload
+def sample(
+    compiled_model: CompiledModel,
+    *,
+    draws: int,
+    tune: int,
+    chains: int,
+    cores: int,
+    seed: Optional[int],
+    save_warmup: bool,
+    progress_bar: bool,
+    init_mean: Optional[np.ndarray],
+    return_raw_trace: bool,
+    blocking: Literal[False],
+    **kwargs,
+) -> _BackgroundSampler: ...
+
+
 def sample(
     compiled_model: CompiledModel,
     *,
@@ -136,7 +335,8 @@ def sample(
     save_warmup: bool = True,
     progress_bar: bool = True,
     init_mean: Optional[np.ndarray] = None,
-    return_raw_trace=False,
+    return_raw_trace: bool = False,
+    blocking: bool = True,
     **kwargs,
 ) -> arviz.InferenceData:
     """Sample the posterior distribution for a compiled model.
@@ -211,49 +411,29 @@ def sample(
     if init_mean is None:
         init_mean = np.zeros(compiled_model.n_dim)
 
-    sampler = compiled_model._make_sampler(settings, init_mean, chains, cores, seed)
+    sampler = _BackgroundSampler(
+        compiled_model,
+        settings,
+        init_mean,
+        chains,
+        cores,
+        seed,
+        draws,
+        tune,
+        progress_bar=progress_bar,
+        save_warmup=save_warmup,
+        return_raw_trace=return_raw_trace,
+    )
+
+    if not blocking:
+        return sampler
 
     try:
-        num_divs = 0
-        chains_tuning = chains
-        bar = fastprogress.progress_bar(
-            sampler,
-            total=chains * (draws + tune),
-            display=progress_bar,
-        )
-        try:
-            for info in bar:
-                if info.draw == tune - 1:
-                    chains_tuning -= 1
-                if info.is_diverging and info.draw > tune:
-                    num_divs += 1
-                bar.comment = (
-                    f" Chains in warmup: {chains_tuning}, Divergences: {num_divs}"
-                )
-        except KeyboardInterrupt:
-            pass
-    finally:
-        results = sampler.finalize()
-
-    dims = {name: list(dim) for name, dim in compiled_model.dims.items()}
-    dims["mass_matrix_inv"] = ["unconstrained_parameter"]
-    dims["gradient"] = ["unconstrained_parameter"]
-    dims["unconstrained_draw"] = ["unconstrained_parameter"]
-    dims["divergence_start"] = ["unconstrained_parameter"]
-    dims["divergence_start_gradient"] = ["unconstrained_parameter"]
-    dims["divergence_end"] = ["unconstrained_parameter"]
-    dims["divergence_momentum"] = ["unconstrained_parameter"]
-
-    if return_raw_trace:
-        return results
-    else:
-        return _trace_to_arviz(
-            results,
-            tune,
-            compiled_model.shapes,
-            dims=dims,
-            coords={
-                name: pd.Index(vals) for name, vals in compiled_model.coords.items()
-            },
-            save_warmup=save_warmup,
-        )
+        result = sampler.wait()
+    except KeyboardInterrupt:
+        result = sampler.abort()
+    except:
+        sampler.cancel()
+        raise
+
+    return result
diff --git a/src/sampler.rs b/src/sampler.rs
index c12bb80..ee0acd3 100644
--- a/src/sampler.rs
+++ b/src/sampler.rs
@@ -1,14 +1,16 @@
 use std::{
     sync::{
         self,
-        mpsc::{channel, sync_channel, RecvTimeoutError, SyncSender},
+        mpsc::{
+            channel, sync_channel, Receiver, RecvError, RecvTimeoutError, SyncSender, TryRecvError,
+        },
         Arc,
     },
-    thread,
+    thread::{self, JoinHandle},
     time::{Duration, Instant},
 };
 
-use anyhow::{Result, Context};
+use anyhow::{anyhow, Context, Result};
 use arrow2::array::Array;
 use itertools::Itertools;
 use nuts_rs::{new_sampler, ArrowBuilder, Chain, CpuLogpFunc, SampleStats, SamplerArgs};
@@ -80,20 +82,23 @@ pub(crate) trait Model: Send + Sync + 'static {
     }
 }
 
+type ChainResult = (u64, Box<dyn Array>, Option<Box<dyn Array>>);
+pub(crate) type SamplerResult = Vec<(Box<dyn Array>, Option<Box<dyn Array>>)>;
+
 pub(crate) struct Sampler {
     /// For each thread we return the chain id, the draws and the stats
-    main_thread: thread::JoinHandle<Result<Vec<(u64, Box<dyn Array>, Option<Box<dyn Array>>)>>>,
+    main_thread: thread::JoinHandle<Result<Vec<ChainResult>>>,
     updates: sync::mpsc::Receiver<Box<dyn SampleStats>>,
 }
 
 impl Sampler {
-    fn run_sampler<M: Model>(
+    fn run_chain<M: Model>(
         seed: Option<u64>,
         chain: u64,
         model: Arc<M>,
         settings: SamplerArgs,
         updates: SyncSender<Box<dyn SampleStats>>,
-    ) -> Result<(u64, Box<dyn Array>, Option<Box<dyn Array>>)> {
+    ) -> Result<ChainResult> {
         let mut rng = if let Some(seed) = seed {
             ChaCha8Rng::seed_from_u64(seed)
         } else {
@@ -151,7 +156,9 @@ impl Sampler {
 
         Ok((
             chain,
-            trace.finalize().context("Failed to finalize the trace object")?,
+            trace
+                .finalize()
+                .context("Failed to finalize the trace object")?,
             stats.finalize().map(|x| x.boxed()),
         ))
     }
@@ -164,13 +171,7 @@ impl Sampler {
         seed: Option<u64>,
     ) -> Self {
         let model = Arc::new(model);
-        let (send_updates, updates) = sync_channel(
-            chains
-                .checked_mul(4)
-                .expect("Invalid number of chains")
-                .try_into()
-                .expect("Invalid number of chains"),
-        );
+        let (send_updates, updates) = sync_channel((4 * chains) as usize);
 
         let thread = thread::spawn(move || {
             let pool = rayon::ThreadPoolBuilder::new()
@@ -188,7 +189,7 @@ impl Sampler {
                 let send_updates = send_updates.clone();
 
                 pool.spawn(move || {
-                    let _ = tx.send(Sampler::run_sampler(
+                    let _ = tx.send(Sampler::run_chain(
                         seed,
                         chain,
                         model,
@@ -208,24 +209,171 @@ impl Sampler {
             updates,
         }
     }
+
+    pub fn finalize(self) -> Result<SamplerResult> {
+        drop(self.updates);
+        match self.main_thread.join() {
+            Ok(vals) => vals.map(|vals| {
+                vals.into_iter()
+                    .sorted_by_key(|&(chain, _, _)| chain)
+                    .map(|(_, tr, stats)| (tr, stats))
+                    .collect()
+            }),
+            Err(panic) => std::panic::resume_unwind(panic),
+        }
+    }
 }
 
-impl Sampler {
-    pub fn next_draw_timeout(
-        &mut self,
-        timeout: Duration,
-    ) -> Result<Box<dyn SampleStats>, RecvTimeoutError> {
-        self.updates.recv_timeout(timeout)
+type SamplerCallback = Box<dyn FnMut(Box<dyn SampleStats>) + Send>;
+
+enum Command {
+    Pause,
+    Resume,
+}
+
+pub(crate) struct SamplerControl {
+    poll_thread: JoinHandle<()>,
+    commands: SyncSender<Command>,
+    results: Receiver<Result<SamplerResult>>,
+}
+
+pub(crate) enum SamplerWaitResult {
+    Result(Result<SamplerResult>),
+    Timeout(SamplerControl),
+}
+
+impl SamplerControl {
+    pub(crate) fn new(sampler: Sampler, mut callback: SamplerCallback) -> Self {
+        let (command_tx, command_rx) = sync_channel(0);
+        let (result_tx, result_rx) = sync_channel(0);
+        let poll_thread = thread::spawn(move || {
+            let mut message = command_rx.try_recv();
+            loop {
+                match message {
+                    Ok(Command::Pause) => {
+                        message = command_rx.recv().map_err(|e| e.into());
+                    }
+                    Ok(Command::Resume) => {
+                        message = command_rx.try_recv();
+                    }
+                    Err(TryRecvError::Empty) => {
+                        if let Ok(draw) = sampler.updates.recv() {
+                            callback(draw);
+                            message = command_rx.try_recv();
+                            continue;
+                        } else {
+                            drop(command_rx);
+                            result_tx
+                                .send(sampler.finalize())
+                                .expect("Could not send results to parent thread");
+                            return;
+                        }
+                    }
+                    Err(TryRecvError::Disconnected) => {
+                        drop(command_rx);
+                        result_tx
+                            .send(sampler.finalize())
+                            .expect("Could not send results to parent thread");
+                        return;
+                    }
+                }
+            }
+        });
+
+        SamplerControl {
+            poll_thread,
+            commands: command_tx,
+            results: result_rx,
+        }
     }
 
-    pub fn finalize(self) -> Result<Vec<(Box<dyn Array>, Option<Box<dyn Array>>)>> {
-        drop(self.updates);
-        let vals = self.main_thread.join().unwrap();
-        vals.map(|vals| {
-            vals.into_iter()
-                .sorted_by_key(|&(chain, _, _)| chain)
-                .map(|(_, tr, stats)| (tr, stats))
-                .collect()
-        })
+    pub(crate) fn pause(&mut self) -> Result<()> {
+        Ok(self
+            .commands
+            .send(Command::Pause)
+            .context("Could not send pause command to sampler")?)
+    }
+
+    pub(crate) fn resume(&mut self) -> Result<()> {
+        Ok(self
+            .commands
+            .send(Command::Resume)
+            .context("Could not send resume command to sampler")?)
+    }
+
+    pub(crate) fn wait_timeout(mut self, timeout: Duration) -> SamplerWaitResult {
+        // We want to ignore this error, because this only
+        // fails if the sampler thread is done, and in this case
+        // we don't need to do anything.
+        let _ = self.resume();
+        match self.results.recv_timeout(timeout) {
+            Err(RecvTimeoutError::Timeout) => SamplerWaitResult::Timeout(self),
+            Err(err @ RecvTimeoutError::Disconnected) => SamplerWaitResult::Result(
+                Err(err).context("Could not get sampler result in wait, sampler thread is dead"),
+            ),
+            Ok(result) => SamplerWaitResult::Result(result),
+        }
+    }
+
+    pub(crate) fn abort(self) -> Result<SamplerResult> {
+        drop(self.commands);
+
+        let result = match self.results.recv() {
+            Err(err @ RecvError) => {
+                Err(err).context("Could not get sampler result in abort, sampler thread is dead")
+            }
+            Ok(result) => result,
+        };
+
+        drop(self.results);
+        let poll_result = self.poll_thread.join();
+        if let Err(_) = poll_result {
+            Err(anyhow!("Sample polling thread paniced."))?
+        }
+        result
+    }
+
+    pub(crate) fn finalize(mut self) -> Result<SamplerResult> {
+        // We want to ignore this error, because this only
+        // fails if the sampler thread is done, and in this case
+        // we don't need to do anything.
+        let _ = self.resume();
+        let result = match self.results.recv() {
+            Err(err @ RecvError) => {
+                Err(err).context("Could not get sampler result in finalize, sampler thread is dead")
+            }
+            Ok(result) => result,
+        };
+
+        drop(self.commands);
+        drop(self.results);
+        let poll_result = self.poll_thread.join();
+        if let Err(_) = poll_result {
+            Err(anyhow!("Sample polling thread paniced."))?
+        }
+        result
+    }
+
+    pub(crate) fn try_finalize(mut self) -> SamplerWaitResult {
+        // We want to ignore this error, because this only
+        // fails if the sampler thread is done, and in this case
+        // we don't need to do anything.
+        let _ = self.resume();
+        let result = match self.results.try_recv() {
+            Err(err @ TryRecvError::Disconnected) => Err(err)
+                .context("Could not get sampler result in try_finalize, sampler thread is dead."),
+            Err(TryRecvError::Empty) => {
+                return SamplerWaitResult::Timeout(self);
+            }
+            Ok(result) => result,
+        };
+
+        drop(self.commands);
+        drop(self.results);
+        let poll_result = self.poll_thread.join();
+        if let Err(_) = poll_result {
+            return SamplerWaitResult::Result(Err(anyhow!("Sample polling thread paniced.")));
+        }
+        SamplerWaitResult::Result(result)
     }
 }
diff --git a/src/stan.rs b/src/stan.rs
index 67a5e79..a8e1606 100644
--- a/src/stan.rs
+++ b/src/stan.rs
@@ -280,13 +280,15 @@ pub struct StanTrace<'model> {
 
 impl<'model> Trace for StanTrace<'model> {
     fn append_value(&mut self, point: &[f64]) -> anyhow::Result<()> {
-        self.inner.param_constrain(
-            point,
-            true,
-            true,
-            &mut self.expanded_buffer,
-            Some(&mut self.rng),
-        ).context("Failed to constrain the parameters of the draw")?;
+        self.inner
+            .param_constrain(
+                point,
+                true,
+                true,
+                &mut self.expanded_buffer,
+                Some(&mut self.rng),
+            )
+            .context("Failed to constrain the parameters of the draw")?;
         for (var, trace) in self.model.variables.iter().zip_eq(self.trace.iter_mut()) {
             let slice = &self.expanded_buffer[var.start_idx..var.end_idx];
             assert!(slice.len() == var.size);
diff --git a/src/wrapper.rs b/src/wrapper.rs
index 37b2370..a98bfdb 100644
--- a/src/wrapper.rs
+++ b/src/wrapper.rs
@@ -1,16 +1,16 @@
-use std::{sync::mpsc::RecvTimeoutError, time::Duration};
+use std::time::{Duration, Instant};
 
 use crate::{
     pymc::{ExpandFunc, LogpFunc, PyMcModel},
-    sampler::Sampler,
+    sampler::{Sampler, SamplerControl, SamplerResult, SamplerWaitResult},
     stan::{StanLibrary, StanModel},
 };
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use arrow2::{array::Array, datatypes::Field};
 use nuts_rs::{SampleStats, SamplerArgs};
 use pyo3::{
-    exceptions::PyValueError,
+    exceptions::{PyTimeoutError, PyValueError},
     ffi::Py_uintptr_t,
     prelude::*,
     types::{PyList, PyTuple},
@@ -212,9 +212,25 @@ impl PySamplerArgs {
     }
 }
 
+pub(crate) enum SamplerState {
+    Running(SamplerControl),
+    Finished(Result<SamplerResult>),
+    Empty,
+}
+
 #[pyclass]
 struct PySampler {
-    sampler: Option<Sampler>,
+    state: SamplerState,
+}
+
+fn make_callback(callback: Option<Py<PyAny>>) -> Box<dyn FnMut(Box<dyn SampleStats>) + Send> {
+    if let Some(callback) = callback {
+        Box::new(move |stats| {
+            let _ = Python::with_gil(|py| callback.call1(py, (PySampleStats { stats },)));
+        })
+    } else {
+        Box::new(|_| {})
+    }
 }
 
 #[pymethods]
@@ -226,9 +242,12 @@ impl PySampler {
         cores: usize,
         model: PyMcModel,
         seed: Option<u64>,
+        callback: Option<Py<PyAny>>,
     ) -> PyResult<PySampler> {
+        let sampler = Sampler::new(model, settings.inner, cores, chains, seed);
+        let control = SamplerControl::new(sampler, make_callback(callback));
         Ok(PySampler {
-            sampler: Some(Sampler::new(model, settings.inner, cores, chains, seed)),
+            state: SamplerState::Running(control),
         })
     }
 
@@ -239,50 +258,147 @@ impl PySampler {
         cores: usize,
         model: StanModel,
         seed: Option<u64>,
+        callback: Option<Py<PyAny>>,
     ) -> PyResult<PySampler> {
+        let sampler = Sampler::new(model, settings.inner, cores, chains, seed);
+        let control = SamplerControl::new(sampler, make_callback(callback));
         Ok(PySampler {
-            sampler: Some(Sampler::new(model, settings.inner, cores, chains, seed)),
+            state: SamplerState::Running(control),
         })
     }
 
-    fn __iter__(self_: PyRef<Self>) -> PyResult<Py<PySampler>> {
-        if self_.sampler.is_none() {
-            Err(PyValueError::new_err("Sampler is finished"))
-        } else {
-            Ok(self_.into())
+    fn is_finished(&mut self) -> PyResult<bool> {
+        let state = std::mem::replace(&mut self.state, SamplerState::Empty);
+
+        let SamplerState::Running(control) = state else {
+            let _ = std::mem::replace(&mut self.state, state);
+            return Ok(true);
+        };
+
+        match control.try_finalize() {
+            SamplerWaitResult::Result(result) => {
+                let _ = std::mem::replace(&mut self.state, SamplerState::Finished(result));
+                Ok(true)
+            }
+            SamplerWaitResult::Timeout(control) => {
+                let _ = std::mem::replace(&mut self.state, SamplerState::Running(control));
+                Ok(false)
+            }
         }
     }
 
-    fn __next__(mut self_: PyRefMut<Self>, py: Python<'_>) -> PyResult<Option<PySampleStats>> {
-        let sampler = self_
-            .sampler
-            .as_mut()
-            .ok_or_else(|| PyValueError::new_err("Sampler is already finished"))?;
-        py.allow_threads(|| loop {
-            match sampler.next_draw_timeout(Duration::from_millis(100)) {
-                Err(RecvTimeoutError::Timeout) => {
-                    Python::with_gil(|py| py.check_signals())?;
-                    continue;
-                }
-                Err(RecvTimeoutError::Disconnected) => {
-                    return Ok(None);
+    fn pause(&mut self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            if let SamplerState::Running(ref mut control) = self.state {
+                control.pause()?
+            }
+            Ok(())
+        })
+    }
+
+    fn resume(&mut self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            if let SamplerState::Running(ref mut control) = self.state {
+                control.resume()?
+            }
+            Ok(())
+        })
+    }
+
+    fn wait(&mut self, py: Python<'_>, timeout_seconds: Option<f64>) -> PyResult<()> {
+        py.allow_threads(|| {
+            let timeout = match timeout_seconds {
+                Some(val) => Some(Duration::try_from_secs_f64(val).context("Invalid timeout")?),
+                None => None,
+            };
+
+            let state = std::mem::replace(&mut self.state, SamplerState::Empty);
+
+            let SamplerState::Running(mut control) = state else {
+                let _ = std::mem::replace(&mut self.state, state);
+                return Ok(());
+            };
+
+            let start_time = Instant::now();
+            let step = Duration::from_millis(100);
+
+            let (final_state, retval) = loop {
+                let time_so_far = Instant::now().saturating_duration_since(start_time);
+                let next_timeout = match timeout {
+                    Some(timeout) => {
+                        let Some(remaining) = timeout.checked_sub(time_so_far) else {
+                            break (
+                                SamplerState::Running(control),
+                                Err(PyTimeoutError::new_err(
+                                    "Timeout while waiting for sampler to finish",
+                                )),
+                            );
+                        };
+                        remaining.min(step)
+                    }
+                    None => step,
+                };
+
+                match control.wait_timeout(next_timeout) {
+                    SamplerWaitResult::Result(result) => {
+                        break (SamplerState::Finished(result), Ok(()))
+                    }
+                    SamplerWaitResult::Timeout(new_control) => {
+                        control = new_control;
+                    }
                 }
-                Ok(stats) => {
-                    let stats = PySampleStats { stats };
-                    return Ok(Some(stats));
+
+                if let Err(err) = Python::with_gil(|py| py.check_signals()) {
+                    break (SamplerState::Running(control), Err(err));
                 }
             };
+
+            let _ = std::mem::replace(&mut self.state, final_state);
+            retval
+        })
+    }
+
+    fn abort(&mut self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            let state = std::mem::replace(&mut self.state, SamplerState::Empty);
+
+            let SamplerState::Running(control) = state else {
+                let _ = std::mem::replace(&mut self.state, state);
+                return Ok(());
+            };
+
+            let result = control.abort();
+            let _ = std::mem::replace(&mut self.state, SamplerState::Finished(result));
+            Ok(())
         })
     }
 
-    fn finalize(&mut self, py: Python<'_>) -> PyResult<Py<PyList>> {
-        let sampler = self
-            .sampler
-            .take()
-            .ok_or_else(|| PyValueError::new_err("Sampler is empty"))?;
-        let values = sampler
-            .finalize()
-            .map_err(|e| PyValueError::new_err(format!("Sampling failed: {:?}", e)))?;
+    fn finalize(&mut self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            let state = std::mem::replace(&mut self.state, SamplerState::Empty);
+
+            let SamplerState::Running(control) = state else {
+                let _ = std::mem::replace(&mut self.state, state);
+                return Ok(());
+            };
+
+            let result = control.finalize();
+            let _ = std::mem::replace(&mut self.state, SamplerState::Finished(result));
+            Ok(())
+        })
+    }
+
+    fn extract_results(&mut self, py: Python<'_>) -> PyResult<Py<PyList>> {
+        let state = std::mem::replace(&mut self.state, SamplerState::Empty);
+
+        let SamplerState::Finished(values) = state else {
+            let _ = std::mem::replace(&mut self.state, state);
+            // Must have called finalize or abort
+            return Err(anyhow::anyhow!("Sampler is not finalized"))?;
+        };
+
+        let values =
+            values.map_err(|e| PyValueError::new_err(format!("Sampling failed: {:?}", e)))?;
 
         let list = PyList::new(
             py,
@@ -305,6 +421,14 @@ impl PySampler {
         );
         Ok(list.into_py(py))
     }
+
+    fn is_empty(&self) -> bool {
+        match self.state {
+            SamplerState::Running(_) => false,
+            SamplerState::Finished(_) => false,
+            SamplerState::Empty => true,
+        }
+    }
 }
 
 fn export_array(py: Python<'_>, name: String, data: Box<dyn Array>) -> PyResult<PyObject> {


Test patch:
diff --git a/tests/test_pymc.py b/tests/test_pymc.py
index 101f1a0..5c443d8 100644
--- a/tests/test_pymc.py
+++ b/tests/test_pymc.py
@@ -15,6 +15,38 @@ def test_pymc_model():
     trace.posterior.a  # noqa: B018
 
 
+def test_blocking():
+    with pm.Model() as model:
+        pm.Normal("a")
+
+    compiled = nutpie.compile_pymc_model(model)
+    sampler = nutpie.sample(compiled, chains=1, blocking=False)
+    trace = sampler.wait()
+    trace.posterior.a  # noqa: B018
+
+
+@pytest.mark.timeout(2)
+def test_wait_timeout():
+    with pm.Model() as model:
+        pm.Normal("a", shape=100_000)
+    compiled = nutpie.compile_pymc_model(model)
+    sampler = nutpie.sample(compiled, chains=1, blocking=False)
+    with pytest.raises(TimeoutError):
+        sampler.wait(timeout=0.1)
+    sampler.cancel()
+
+
+@pytest.mark.timeout(2)
+def test_pause():
+    with pm.Model() as model:
+        pm.Normal("a", shape=100_000)
+    compiled = nutpie.compile_pymc_model(model)
+    sampler = nutpie.sample(compiled, chains=1, blocking=False)
+    sampler.pause()
+    sampler.resume()
+    sampler.cancel()
+
+
 def test_pymc_model_with_coordinate():
     with pm.Model() as model:
         model.add_coord("foo", length=5)