diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..8f61a8e
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,2 @@
+# SCM syntax highlighting
+pixi.lock linguist-language=YAML linguist-generated=true
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 5ab1d1b..41609a5 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -31,16 +31,32 @@ jobs:
             target: aarch64
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
       - name: Install uv
         uses: astral-sh/setup-uv@v5
+      - uses: actions/setup-python@v5
+        with:
+          python-version: |
+            3.10
+            3.11
+            3.12
+            3.13
+      # Leave out 3.13 on aarch due to an issue in pyo3/rust-numpy 0.23.4
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        if: ${{ matrix.platform.target == 'aarch64' }}
+        with:
+          target: ${{ matrix.platform.target }}
+          args: --release --out dist --interpreter 3.10 3.11 3.12 --zig
+          sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
+          manylinux: auto
+          before-script-linux: |
+            dnf install -y clang-libs clang || sudo apt install llvm-dev libclang-dev clang
       - name: Build wheels
         uses: PyO3/maturin-action@v1
+        if: ${{ matrix.platform.target == 'x86_64' }}
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --find-interpreter --zig
+          args: --release --out dist --interpreter 3.10 3.11 3.12 3.13 --zig
           sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
           manylinux: auto
           before-script-linux: |
@@ -160,7 +176,11 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: |
+            3.10
+            3.11
+            3.12
+          # 3.13 leave out 3.13 due to a segfault
           architecture: ${{ matrix.platform.target }}
       - name: Install uv
         uses: astral-sh/setup-uv@v5
@@ -211,7 +231,11 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: |
+            3.10
+            3.11
+            3.12
+            3.13
       - name: Install uv
         uses: astral-sh/setup-uv@v5
       - uses: maxim-lobanov/setup-xcode@v1
diff --git a/.gitignore b/.gitignore
index f0721f7..088a555 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,7 +8,18 @@ tvm_libs/*
 notebooks/*.stan
 notebooks/*.csv
 notebooks/*.hpp
+notebooks/radon*
 perf.data*
 wheels
 .vscode/
 *~
+.zed
+.cargo
+*traces*
+.pyrightconfig.json
+*.zarr
+book
+docs/_site
+.quarto
+example-iree
+posteriordb
diff --git a/Cargo.lock b/Cargo.lock
index d6363e2..1107cae 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2,6 +2,23 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "adler2"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
+
+[[package]]
+name = "aes"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
+dependencies = [
+ "cfg-if",
+ "cipher",
+ "cpufeatures",
+]
+
 [[package]]
 name = "ahash"
 version = "0.8.11"
@@ -10,10 +27,10 @@ checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
 dependencies = [
  "cfg-if",
  "const-random",
- "getrandom",
+ "getrandom 0.2.15",
  "once_cell",
  "version_check",
- "zerocopy",
+ "zerocopy 0.7.35",
 ]
 
 [[package]]
@@ -60,9 +77,9 @@ checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04"
 
 [[package]]
 name = "arrow"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "05048a8932648b63f21c37d88b552ccc8a65afb6dfe9fc9f30ce79174c2e7a85"
+checksum = "6422e12ac345a0678d7a17e316238e3a40547ae7f92052b77bd86d5e0239f3fc"
 dependencies = [
  "arrow-arith",
  "arrow-array",
@@ -78,24 +95,23 @@ dependencies = [
 
 [[package]]
 name = "arrow-arith"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d8a57966e43bfe9a3277984a14c24ec617ad874e4c0e1d2a1b083a39cfbf22c"
+checksum = "23cf34bb1f48c41d3475927bcc7be498665b8e80b379b88f62a840337f8b8248"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "chrono",
- "half",
  "num",
 ]
 
 [[package]]
 name = "arrow-array"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "16f4a9468c882dc66862cef4e1fd8423d47e67972377d85d80e022786427768c"
+checksum = "fb4a06d507f54b70a277be22a127c8ffe0cec6cd98c0ad8a48e77779bbda8223"
 dependencies = [
  "ahash",
  "arrow-buffer",
@@ -109,9 +125,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-buffer"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c975484888fc95ec4a632cdc98be39c085b1bb518531b0c80c5d462063e5daa1"
+checksum = "d69d326d5ad1cb82dcefa9ede3fee8fdca98f9982756b16f9cb142f4aa6edc89"
 dependencies = [
  "bytes",
  "half",
@@ -120,9 +136,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-cast"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "da26719e76b81d8bc3faad1d4dbdc1bcc10d14704e63dc17fc9f3e7e1e567c8e"
+checksum = "626e65bd42636a84a238bed49d09c8777e3d825bf81f5087a70111c2831d9870"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -140,9 +156,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-data"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dd9d6f18c65ef7a2573ab498c374d8ae364b4a4edf67105357491c031f716ca5"
+checksum = "1858e7c7d01c44cf71c21a85534fd1a54501e8d60d1195d0d6fbcc00f4b10754"
 dependencies = [
  "arrow-buffer",
  "arrow-schema",
@@ -152,26 +168,23 @@ dependencies = [
 
 [[package]]
 name = "arrow-ord"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "42745f86b1ab99ef96d1c0bcf49180848a64fe2c7a7a0d945bc64fa2b21ba9bc"
+checksum = "8aa06e5f267dc53efbacb933485c79b6fc1685d3ffbe870a16ce4e696fb429da"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "arrow-select",
- "half",
- "num",
 ]
 
 [[package]]
 name = "arrow-row"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4cd09a518c602a55bd406bcc291a967b284cfa7a63edfbf8b897ea4748aad23c"
+checksum = "66f1144bb456a2f9d82677bd3abcea019217e572fc8f07de5a7bac4b2c56eb2c"
 dependencies = [
- "ahash",
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
@@ -181,18 +194,18 @@ dependencies = [
 
 [[package]]
 name = "arrow-schema"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9e972cd1ff4a4ccd22f86d3e53e835c2ed92e0eea6a3e8eadb72b4f1ac802cf8"
+checksum = "105f01ec0090259e9a33a9263ec18ff223ab91a0ea9fbc18042f7e38005142f6"
 dependencies = [
  "bitflags",
 ]
 
 [[package]]
 name = "arrow-select"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "600bae05d43483d216fb3494f8c32fdbefd8aa4e1de237e790dbb3d9f44690a3"
+checksum = "f690752fdbd2dee278b5f1636fefad8f2f7134c85e20fd59c4199e15a39a6807"
 dependencies = [
  "ahash",
  "arrow-array",
@@ -204,9 +217,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-string"
-version = "52.2.0"
+version = "54.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f0dc1985b67cb45f6606a248ac2b4a288849f196bab8c657ea5589f47cdd55e6"
+checksum = "d0fff9cd745a7039b66c47ecaf5954460f9fa12eed628f65170117ea93e64ee0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -240,6 +253,12 @@ version = "0.22.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
 
+[[package]]
+name = "base64ct"
+version = "1.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b"
+
 [[package]]
 name = "bindgen"
 version = "0.71.1"
@@ -255,9 +274,9 @@ dependencies = [
  "proc-macro2",
  "quote",
  "regex",
- "rustc-hash 2.1.1",
+ "rustc-hash",
  "shlex",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -311,7 +330,7 @@ checksum = "3fa76293b4f7bb636ab88fd78228235b5248b4d05cc589aed610f954af5d7c7a"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -326,6 +345,27 @@ version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9"
 
+[[package]]
+name = "bzip2"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8"
+dependencies = [
+ "bzip2-sys",
+ "libc",
+]
+
+[[package]]
+name = "bzip2-sys"
+version = "0.1.12+1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72ebc2f1a417f01e1da30ef264ee86ae31d2dcd2d603ea283d3c244a883ca2a9"
+dependencies = [
+ "cc",
+ "libc",
+ "pkg-config",
+]
+
 [[package]]
 name = "cast"
 version = "0.3.0"
@@ -338,6 +378,8 @@ version = "1.2.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c7777341816418c02e033934a09f20dc0ccaf65a5201ef8a450ae0105a573fda"
 dependencies = [
+ "jobserver",
+ "libc",
  "shlex",
 ]
 
@@ -395,6 +437,16 @@ dependencies = [
  "half",
 ]
 
+[[package]]
+name = "cipher"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
+dependencies = [
+ "crypto-common",
+ "inout",
+]
+
 [[package]]
 name = "clang-sys"
 version = "1.8.1"
@@ -408,18 +460,18 @@ dependencies = [
 
 [[package]]
 name = "clap"
-version = "4.5.28"
+version = "4.5.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3e77c3243bd94243c03672cb5154667347c457ca271254724f9f393aee1c05ff"
+checksum = "8acebd8ad879283633b343856142139f2da2317c96b05b4dd6181c61e2480184"
 dependencies = [
  "clap_builder",
 ]
 
 [[package]]
 name = "clap_builder"
-version = "4.5.27"
+version = "4.5.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1b26884eb4b57140e4d2d93652abfa49498b938b3c9179f9fc487b0acc3edad7"
+checksum = "f6ba32cbda51c7e1dfd49acc1457ba1a7dec5b64fe360e828acb13ca8dc9c2f9"
 dependencies = [
  "anstyle",
  "clap_lex",
@@ -431,12 +483,6 @@ version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6"
 
-[[package]]
-name = "coe-rs"
-version = "0.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e8f1e641542c07631228b1e0dc04b69ae3c1d58ef65d5691a439711d805c698"
-
 [[package]]
 name = "console"
 version = "0.15.10"
@@ -465,11 +511,17 @@ version = "0.1.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
 dependencies = [
- "getrandom",
+ "getrandom 0.2.15",
  "once_cell",
  "tiny-keccak",
 ]
 
+[[package]]
+name = "constant_time_eq"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
+
 [[package]]
 name = "core-foundation-sys"
 version = "0.8.7"
@@ -485,6 +537,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "crc32fast"
+version = "1.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "criterion"
 version = "0.5.1"
@@ -563,10 +624,13 @@ dependencies = [
 ]
 
 [[package]]
-name = "dbgf"
-version = "0.1.2"
+name = "deranged"
+version = "0.3.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e6ca96b45ca70b8045e0462f191bd209fcb3c3bfe8dbfb1257ada54c4dd59169"
+checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4"
+dependencies = [
+ "powerfmt",
+]
 
 [[package]]
 name = "digest"
@@ -576,16 +640,7 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
 dependencies = [
  "block-buffer",
  "crypto-common",
-]
-
-[[package]]
-name = "dyn-stack"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "56e53799688f5632f364f8fb387488dd05db9fe45db7011be066fc20e7027f8b"
-dependencies = [
- "bytemuck",
- "reborrow",
+ "subtle",
 ]
 
 [[package]]
@@ -615,19 +670,10 @@ version = "0.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc"
 dependencies = [
- "heck 0.5.0",
+ "heck",
  "proc-macro2",
  "quote",
- "syn 2.0.98",
-]
-
-[[package]]
-name = "equator"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c35da53b5a021d2484a7cc49b2ac7f2d840f8236a286f84202369bd338d761ea"
-dependencies = [
- "equator-macro 0.2.1",
+ "syn",
 ]
 
 [[package]]
@@ -636,18 +682,7 @@ version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc"
 dependencies = [
- "equator-macro 0.4.2",
-]
-
-[[package]]
-name = "equator-macro"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3bf679796c0322556351f287a51b49e48f7c4986e727b5dd78c972d30e2e16cc"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.98",
+ "equator-macro",
 ]
 
 [[package]]
@@ -658,59 +693,75 @@ checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
 name = "faer"
-version = "0.19.4"
+version = "0.21.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "64bc4855cb2792ae3520e8af22051a47a6d6dc8300ebc0ddf51ad73f65bd0dc9"
+checksum = "4fe562e9ab13fe9d1249786e560a8f52237f6ab7520bdda9a022997fb0c081fe"
 dependencies = [
  "bytemuck",
- "coe-rs",
- "dbgf",
- "dyn-stack 0.10.0",
- "equator 0.4.2",
- "faer-entity",
+ "dyn-stack",
+ "equator",
+ "faer-macros",
+ "faer-traits",
  "gemm",
+ "generativity",
  "libm",
- "matrixcompare",
- "matrixcompare-core",
- "nano-gemm",
  "npyz",
  "num-complex",
  "num-traits",
- "paste",
- "rand",
- "rand_distr",
- "rayon",
+ "pulp",
  "reborrow",
- "serde",
 ]
 
 [[package]]
-name = "faer-entity"
-version = "0.19.2"
+name = "faer-macros"
+version = "0.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c9c752ab2bff6f0b9597c6a1adc0112f7fd41fb343bc5a009a6274ae9d32fd03"
+checksum = "9d0a255d1442b5825c61812a7eafda9034ec53d969c98555251085e148428e6a"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "faer-traits"
+version = "0.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2d0172aefb5f869561e558d5390657f1aa98ca3c51a09be69a4687064ebfb9a"
 dependencies = [
  "bytemuck",
- "coe-rs",
+ "dyn-stack",
+ "faer-macros",
+ "generativity",
  "libm",
  "num-complex",
  "num-traits",
- "pulp 0.18.22",
+ "pulp",
  "reborrow",
 ]
 
+[[package]]
+name = "flate2"
+version = "1.0.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c936bfdafb507ebbf50b8074c54fa31c5be9a1e7e5f467dd659697041407d07c"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide",
+]
+
 [[package]]
 name = "gemm"
 version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-c32",
  "gemm-c64",
  "gemm-common",
@@ -730,7 +781,7 @@ version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-common",
  "num-complex",
  "num-traits",
@@ -745,7 +796,7 @@ version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-common",
  "num-complex",
  "num-traits",
@@ -761,16 +812,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3"
 dependencies = [
  "bytemuck",
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "half",
  "libm",
  "num-complex",
  "num-traits",
  "once_cell",
  "paste",
- "pulp 0.21.4",
+ "pulp",
  "raw-cpuid",
- "rayon",
  "seq-macro",
  "sysctl",
 ]
@@ -781,7 +831,7 @@ version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-common",
  "gemm-f32",
  "half",
@@ -789,7 +839,6 @@ dependencies = [
  "num-traits",
  "paste",
  "raw-cpuid",
- "rayon",
  "seq-macro",
 ]
 
@@ -799,7 +848,7 @@ version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-common",
  "num-complex",
  "num-traits",
@@ -814,7 +863,7 @@ version = "0.18.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd"
 dependencies = [
- "dyn-stack 0.13.0",
+ "dyn-stack",
  "gemm-common",
  "num-complex",
  "num-traits",
@@ -823,6 +872,12 @@ dependencies = [
  "seq-macro",
 ]
 
+[[package]]
+name = "generativity"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5881e4c3c2433fe4905bb19cfd2b5d49d4248274862b68c27c33d9ba4e13f9ec"
+
 [[package]]
 name = "generic-array"
 version = "0.14.7"
@@ -841,7 +896,19 @@ checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
 dependencies = [
  "cfg-if",
  "libc",
- "wasi",
+ "wasi 0.11.0+wasi-snapshot-preview1",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi 0.13.3+wasi-0.2.2",
+ "windows-targets",
 ]
 
 [[package]]
@@ -864,15 +931,9 @@ dependencies = [
 
 [[package]]
 name = "hashbrown"
-version = "0.14.5"
+version = "0.15.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
-
-[[package]]
-name = "heck"
-version = "0.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
+checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
 
 [[package]]
 name = "heck"
@@ -886,6 +947,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
 
+[[package]]
+name = "hmac"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
+dependencies = [
+ "digest",
+]
+
 [[package]]
 name = "iana-time-zone"
 version = "0.1.61"
@@ -928,6 +998,15 @@ version = "2.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"
 
+[[package]]
+name = "inout"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5"
+dependencies = [
+ "generic-array",
+]
+
 [[package]]
 name = "is-terminal"
 version = "0.4.15"
@@ -972,6 +1051,15 @@ version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674"
 
+[[package]]
+name = "jobserver"
+version = "0.1.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "js-sys"
 version = "0.3.77"
@@ -982,11 +1070,17 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
 [[package]]
 name = "lexical-core"
-version = "0.8.5"
+version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2cde5de06e8d4c2faabc400238f9ae1c74d5412d03a7bd067645ccbc47070e46"
+checksum = "b765c31809609075565a70b4b71402281283aeda7ecaf4818ac14a7b2ade8958"
 dependencies = [
  "lexical-parse-float",
  "lexical-parse-integer",
@@ -997,9 +1091,9 @@ dependencies = [
 
 [[package]]
 name = "lexical-parse-float"
-version = "0.8.5"
+version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "683b3a5ebd0130b8fb52ba0bdc718cc56815b6a097e28ae5a6997d0ad17dc05f"
+checksum = "de6f9cb01fb0b08060209a057c048fcbab8717b4c1ecd2eac66ebfe39a65b0f2"
 dependencies = [
  "lexical-parse-integer",
  "lexical-util",
@@ -1008,9 +1102,9 @@ dependencies = [
 
 [[package]]
 name = "lexical-parse-integer"
-version = "0.8.6"
+version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6d0994485ed0c312f6d965766754ea177d07f9c00c9b82a5ee62ed5b47945ee9"
+checksum = "72207aae22fc0a121ba7b6d479e42cbfea549af1479c3f3a4f12c70dd66df12e"
 dependencies = [
  "lexical-util",
  "static_assertions",
@@ -1018,18 +1112,18 @@ dependencies = [
 
 [[package]]
 name = "lexical-util"
-version = "0.8.5"
+version = "1.0.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5255b9ff16ff898710eb9eb63cb39248ea8a5bb036bea8085b1a767ff6c4e3fc"
+checksum = "5a82e24bf537fd24c177ffbbdc6ebcc8d54732c35b50a3f28cc3f4e4c949a0b3"
 dependencies = [
  "static_assertions",
 ]
 
 [[package]]
 name = "lexical-write-float"
-version = "0.8.5"
+version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "accabaa1c4581f05a3923d1b4cfd124c329352288b7b9da09e766b0668116862"
+checksum = "c5afc668a27f460fb45a81a757b6bf2f43c2d7e30cb5a2dcd3abf294c78d62bd"
 dependencies = [
  "lexical-util",
  "lexical-write-integer",
@@ -1038,9 +1132,9 @@ dependencies = [
 
 [[package]]
 name = "lexical-write-integer"
-version = "0.8.5"
+version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e1b6f3d1f4422866b68192d62f77bc5c700bee84f3069f2469d7bc8c77852446"
+checksum = "629ddff1a914a836fb245616a7888b62903aae58fa771e1d83943035efa0f978"
 dependencies = [
  "lexical-util",
  "static_assertions",
@@ -1068,38 +1162,12 @@ version = "0.2.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"
 
-[[package]]
-name = "lock_api"
-version = "0.4.12"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17"
-dependencies = [
- "autocfg",
- "scopeguard",
-]
-
 [[package]]
 name = "log"
 version = "0.4.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f"
 
-[[package]]
-name = "matrixcompare"
-version = "0.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37832ba820e47c93d66b4360198dccb004b43c74abc3ac1ce1fed54e65a80445"
-dependencies = [
- "matrixcompare-core",
- "num-traits",
-]
-
-[[package]]
-name = "matrixcompare-core"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0bdabb30db18805d5290b3da7ceaccbddba795620b86c02145d688e04900a73"
-
 [[package]]
 name = "matrixmultiply"
 version = "0.3.9"
@@ -1131,11 +1199,20 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
 
+[[package]]
+name = "miniz_oxide"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3b1c9bd4fe1f0f8b387f6eb9eb3b4a1aa26185e5750efb9140301703f62cd1b"
+dependencies = [
+ "adler2",
+]
+
 [[package]]
 name = "multiversion"
-version = "0.7.4"
+version = "0.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4851161a11d3ad0bf9402d90ffc3967bf231768bfd7aeb61755ad06dbf1a142"
+checksum = "7edb7f0ff51249dfda9ab96b5823695e15a052dc15074c9dbf3d118afaf2c201"
 dependencies = [
  "multiversion-macros",
  "target-features",
@@ -1143,96 +1220,28 @@ dependencies = [
 
 [[package]]
 name = "multiversion-macros"
-version = "0.7.4"
+version = "0.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "79a74ddee9e0c27d2578323c13905793e91622148f138ba29738f9dddb835e90"
+checksum = "b093064383341eb3271f42e381cb8f10a01459478446953953c75d24bd339fc0"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 1.0.109",
+ "syn",
  "target-features",
 ]
 
-[[package]]
-name = "nano-gemm"
-version = "0.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f563548d38f390ef9893e4883ec38c1fb312f569e98d76bededdd91a3b41a043"
-dependencies = [
- "equator 0.2.2",
- "nano-gemm-c32",
- "nano-gemm-c64",
- "nano-gemm-codegen",
- "nano-gemm-core",
- "nano-gemm-f32",
- "nano-gemm-f64",
- "num-complex",
-]
-
-[[package]]
-name = "nano-gemm-c32"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a40449e57a5713464c3a1208c4c3301c8d29ee1344711822cf022bc91373a91b"
-dependencies = [
- "nano-gemm-codegen",
- "nano-gemm-core",
- "num-complex",
-]
-
-[[package]]
-name = "nano-gemm-c64"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "743a6e6211358fba85d1009616751e4107da86f4c95b24e684ce85f25c25b3bf"
-dependencies = [
- "nano-gemm-codegen",
- "nano-gemm-core",
- "num-complex",
-]
-
-[[package]]
-name = "nano-gemm-codegen"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "963bf7c7110d55430169dc74c67096375491ed580cd2ef84842550ac72e781fa"
-
-[[package]]
-name = "nano-gemm-core"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fe3fc4f83ae8861bad79dc3c016bd6b0220da5f9de302e07d3112d16efc24aa6"
-
-[[package]]
-name = "nano-gemm-f32"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4e3681b7ce35658f79da94b7f62c60a005e29c373c7111ed070e3bf64546a8bb"
-dependencies = [
- "nano-gemm-codegen",
- "nano-gemm-core",
-]
-
-[[package]]
-name = "nano-gemm-f64"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bc1e619ed04d801809e1f63e61b669d380c4119e8b0cdd6ed184c6b111f046d8"
-dependencies = [
- "nano-gemm-codegen",
- "nano-gemm-core",
-]
-
 [[package]]
 name = "ndarray"
-version = "0.15.6"
+version = "0.16.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
 dependencies = [
  "matrixmultiply",
  "num-complex",
  "num-integer",
  "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
  "rawpointer",
 ]
 
@@ -1289,9 +1298,14 @@ checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
 dependencies = [
  "bytemuck",
  "num-traits",
- "rand",
 ]
 
+[[package]]
+name = "num-conv"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9"
+
 [[package]]
 name = "num-integer"
 version = "0.1.46"
@@ -1341,9 +1355,9 @@ checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
 
 [[package]]
 name = "numpy"
-version = "0.21.0"
+version = "0.23.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec170733ca37175f5d75a5bea5911d6ff45d2cd52849ce98b685394e4f2f37f4"
+checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78"
 dependencies = [
  "libc",
  "ndarray",
@@ -1351,12 +1365,12 @@ dependencies = [
  "num-integer",
  "num-traits",
  "pyo3",
- "rustc-hash 1.1.0",
+ "rustc-hash",
 ]
 
 [[package]]
 name = "nutpie"
-version = "0.13.3"
+version = "0.14.0"
 dependencies = [
  "anyhow",
  "arrow",
@@ -1367,33 +1381,34 @@ dependencies = [
  "numpy",
  "nuts-rs",
  "pyo3",
- "rand",
- "rand_chacha",
+ "rand 0.9.0",
+ "rand_chacha 0.9.0",
  "rand_distr",
  "rayon",
  "smallvec",
- "thiserror 1.0.69",
+ "tch",
+ "thiserror 2.0.11",
  "time-humanize",
  "upon",
 ]
 
 [[package]]
 name = "nuts-rs"
-version = "0.12.1"
+version = "0.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d8573e3b5c83e8ec0570ebbd75dd6fdc7dfcfa5da9b5f9d9d63fedefebbd9cf8"
+checksum = "10e87924d332fce1202087bc67db7ed8f7ef9229da5ec74a5130568f5b7f6ac7"
 dependencies = [
  "anyhow",
  "arrow",
  "faer",
- "itertools 0.13.0",
+ "itertools 0.14.0",
  "multiversion",
- "pulp 0.18.22",
- "rand",
- "rand_chacha",
+ "pulp",
+ "rand 0.9.0",
+ "rand_chacha 0.9.0",
  "rand_distr",
  "rayon",
- "thiserror 1.0.69",
+ "thiserror 2.0.11",
 ]
 
 [[package]]
@@ -1409,26 +1424,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b410bbe7e14ab526a0e86877eb47c6996a2bd7746f027ba551028c925390e4e9"
 
 [[package]]
-name = "parking_lot"
-version = "0.12.3"
+name = "password-hash"
+version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27"
+checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700"
 dependencies = [
- "lock_api",
- "parking_lot_core",
-]
-
-[[package]]
-name = "parking_lot_core"
-version = "0.9.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8"
-dependencies = [
- "cfg-if",
- "libc",
- "redox_syscall",
- "smallvec",
- "windows-targets",
+ "base64ct",
+ "rand_core 0.6.4",
+ "subtle",
 ]
 
 [[package]]
@@ -1455,6 +1458,18 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "pbkdf2"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917"
+dependencies = [
+ "digest",
+ "hmac",
+ "password-hash",
+ "sha2",
+]
+
 [[package]]
 name = "pest"
 version = "2.7.15"
@@ -1486,7 +1501,7 @@ dependencies = [
  "pest_meta",
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -1500,6 +1515,12 @@ dependencies = [
  "sha2",
 ]
 
+[[package]]
+name = "pkg-config"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2"
+
 [[package]]
 name = "plotters"
 version = "0.3.7"
@@ -1534,13 +1555,28 @@ version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
 
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "powerfmt"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04"
 dependencies = [
- "zerocopy",
+ "zerocopy 0.7.35",
 ]
 
 [[package]]
@@ -1550,7 +1586,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6924ced06e1f7dfe3fa48d57b9f74f55d8915f5036121bef647ef4b204895fac"
 dependencies = [
  "proc-macro2",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -1562,18 +1598,6 @@ dependencies = [
  "unicode-ident",
 ]
 
-[[package]]
-name = "pulp"
-version = "0.18.22"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a0a01a0dc67cf4558d279f0c25b0962bd08fc6dec0137699eae304103e882fe6"
-dependencies = [
- "bytemuck",
- "libm",
- "num-complex",
- "reborrow",
-]
-
 [[package]]
 name = "pulp"
 version = "0.21.4"
@@ -1603,16 +1627,16 @@ dependencies = [
 
 [[package]]
 name = "pyo3"
-version = "0.21.2"
+version = "0.23.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a5e00b96a521718e08e03b1a622f01c8a8deb50719335de3f60b3b3950f069d8"
+checksum = "57fe09249128b3173d092de9523eaa75136bf7ba85e0d69eca241c7939c933cc"
 dependencies = [
  "anyhow",
  "cfg-if",
  "indoc",
  "libc",
  "memoffset",
- "parking_lot",
+ "once_cell",
  "portable-atomic",
  "pyo3-build-config",
  "pyo3-ffi",
@@ -1622,9 +1646,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.21.2"
+version = "0.23.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7883df5835fafdad87c0d888b266c8ec0f4c9ca48a5bed6bbb592e8dedee1b50"
+checksum = "1cd3927b5a78757a0d71aa9dff669f903b1eb64b54142a9bd9f757f8fde65fd7"
 dependencies = [
  "once_cell",
  "target-lexicon",
@@ -1632,9 +1656,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.21.2"
+version = "0.23.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "01be5843dc60b916ab4dad1dca6d20b9b4e6ddc8e15f50c47fe6d85f1fb97403"
+checksum = "dab6bb2102bd8f991e7749f130a70d05dd557613e39ed2deeee8e9ca0c4d548d"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -1642,27 +1666,27 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.21.2"
+version = "0.23.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77b34069fc0682e11b31dbd10321cbf94808394c56fd996796ce45217dfac53c"
+checksum = "91871864b353fd5ffcb3f91f2f703a22a9797c91b9ab497b1acac7b07ae509c7"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.21.2"
+version = "0.23.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "08260721f32db5e1a5beae69a55553f56b99bd0e1c3e6e0a5e8851a9d0f5a85c"
+checksum = "43abc3b80bc20f3facd86cd3c60beed58c3e2aa26213f3cda368de39c60a27e4"
 dependencies = [
- "heck 0.4.1",
+ "heck",
  "proc-macro2",
  "pyo3-build-config",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -1681,8 +1705,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
 dependencies = [
  "libc",
- "rand_chacha",
- "rand_core",
+ "rand_chacha 0.3.1",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94"
+dependencies = [
+ "rand_chacha 0.9.0",
+ "rand_core 0.9.0",
+ "zerocopy 0.8.17",
 ]
 
 [[package]]
@@ -1692,7 +1727,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
 dependencies = [
  "ppv-lite86",
- "rand_core",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.9.0",
 ]
 
 [[package]]
@@ -1701,17 +1746,27 @@ version = "0.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
 dependencies = [
- "getrandom",
+ "getrandom 0.2.15",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff"
+dependencies = [
+ "getrandom 0.3.1",
+ "zerocopy 0.8.17",
 ]
 
 [[package]]
 name = "rand_distr"
-version = "0.4.3"
+version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31"
+checksum = "ddc3b5afe4c995c44540865b8ca5c52e6a59fa362da96c5d30886930ddc8da1c"
 dependencies = [
  "num-traits",
- "rand",
+ "rand 0.9.0",
 ]
 
 [[package]]
@@ -1755,15 +1810,6 @@ version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430"
 
-[[package]]
-name = "redox_syscall"
-version = "0.5.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "03a862b389f93e68874fbf580b9de08dd02facb9a788ebadaf4a3fd33cf58834"
-dependencies = [
- "bitflags",
-]
-
 [[package]]
 name = "regex"
 version = "1.11.1"
@@ -1793,12 +1839,6 @@ version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
 
-[[package]]
-name = "rustc-hash"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
-
 [[package]]
 name = "rustc-hash"
 version = "2.1.1"
@@ -1817,6 +1857,16 @@ version = "1.0.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd"
 
+[[package]]
+name = "safetensors"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d93279b86b3de76f820a8854dd06cbc33cfa57a417b19c47f6a25280112fb1df"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "same-file"
 version = "1.0.6"
@@ -1826,12 +1876,6 @@ dependencies = [
  "winapi-util",
 ]
 
-[[package]]
-name = "scopeguard"
-version = "1.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
-
 [[package]]
 name = "seq-macro"
 version = "0.3.5"
@@ -1855,7 +1899,7 @@ checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -1870,6 +1914,17 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "sha1"
+version = "0.10.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
 [[package]]
 name = "sha2"
 version = "0.10.8"
@@ -1900,15 +1955,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
 
 [[package]]
-name = "syn"
-version = "1.0.109"
+name = "subtle"
+version = "2.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
-dependencies = [
- "proc-macro2",
- "quote",
- "unicode-ident",
-]
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
 
 [[package]]
 name = "syn"
@@ -1947,6 +1997,23 @@ version = "0.12.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
 
+[[package]]
+name = "tch"
+version = "0.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa1ed622c8f13b0c42f8b1afa0e5e9ccccd82ecb6c0e904120722ab52fdc5234"
+dependencies = [
+ "half",
+ "lazy_static",
+ "libc",
+ "ndarray",
+ "rand 0.8.5",
+ "safetensors",
+ "thiserror 1.0.69",
+ "torch-sys",
+ "zip",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.69"
@@ -1973,7 +2040,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
 ]
 
 [[package]]
@@ -1984,9 +2051,28 @@ checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
+]
+
+[[package]]
+name = "time"
+version = "0.3.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35e7868883861bd0e56d9ac6efcaaca0d6d5d82a2a7ec8209ff492c07cf37b21"
+dependencies = [
+ "deranged",
+ "num-conv",
+ "powerfmt",
+ "serde",
+ "time-core",
 ]
 
+[[package]]
+name = "time-core"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3"
+
 [[package]]
 name = "time-humanize"
 version = "0.1.3"
@@ -2012,6 +2098,18 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "torch-sys"
+version = "0.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef14f5d239e3d60f4919f536a5dfe1d4f71b27b7abf6fe6875fd3a4b22c2dcd5"
+dependencies = [
+ "anyhow",
+ "cc",
+ "libc",
+ "zip",
+]
+
 [[package]]
 name = "typenum"
 version = "1.17.0"
@@ -2044,9 +2142,9 @@ checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce"
 
 [[package]]
 name = "upon"
-version = "0.8.1"
+version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9fe29601d1624f104fa9a35ea71a5f523dd8bd1cfc8c31f8124ad2b829f013c0"
+checksum = "cc1243af2969e332d5b9b99087eddd44d04a41da8630ed53e06df497b7f5c747"
 
 [[package]]
 name = "version_check"
@@ -2070,6 +2168,15 @@ version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
+[[package]]
+name = "wasi"
+version = "0.13.3+wasi-0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2"
+dependencies = [
+ "wit-bindgen-rt",
+]
+
 [[package]]
 name = "wasm-bindgen"
 version = "0.2.100"
@@ -2092,7 +2199,7 @@ dependencies = [
  "log",
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
  "wasm-bindgen-shared",
 ]
 
@@ -2114,7 +2221,7 @@ checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
  "wasm-bindgen-backend",
  "wasm-bindgen-shared",
 ]
@@ -2239,6 +2346,15 @@ version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
 
+[[package]]
+name = "wit-bindgen-rt"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c"
+dependencies = [
+ "bitflags",
+]
+
 [[package]]
 name = "zerocopy"
 version = "0.7.35"
@@ -2246,7 +2362,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
 dependencies = [
  "byteorder",
- "zerocopy-derive",
+ "zerocopy-derive 0.7.35",
+]
+
+[[package]]
+name = "zerocopy"
+version = "0.8.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa91407dacce3a68c56de03abe2760159582b846c6a4acd2f456618087f12713"
+dependencies = [
+ "zerocopy-derive 0.8.17",
 ]
 
 [[package]]
@@ -2257,5 +2382,65 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.98",
+ "syn",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06718a168365cad3d5ff0bb133aad346959a2074bd4a85c121255a11304a8626"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "zip"
+version = "0.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
+dependencies = [
+ "aes",
+ "byteorder",
+ "bzip2",
+ "constant_time_eq",
+ "crc32fast",
+ "crossbeam-utils",
+ "flate2",
+ "hmac",
+ "pbkdf2",
+ "sha1",
+ "time",
+ "zstd",
+]
+
+[[package]]
+name = "zstd"
+version = "0.11.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "5.0.2+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db"
+dependencies = [
+ "libc",
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.13+zstd.1.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38ff0f21cfee8f97d94cef41359e0c89aa6113028ab0291aa8ca0038995a95aa"
+dependencies = [
+ "cc",
+ "pkg-config",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 16f7ac9..587aa2f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "nutpie"
-version = "0.13.3"
+version = "0.14.0"
 authors = [
   "Adrian Seyboldt <adrian.seyboldt@gmail.com>",
   "PyMC Developers <pymc.devs@gmail.com>",
@@ -22,25 +22,26 @@ name = "_lib"
 crate-type = ["cdylib"]
 
 [dependencies]
-nuts-rs = "0.12.1"
-numpy = "0.21.0"
-rand = "0.8.5"
-thiserror = "1.0.44"
-rand_chacha = "0.3.1"
-rayon = "1.9.0"
+nuts-rs = "0.15.0"
+numpy = "0.23.0"
+rand = "0.9.0"
+thiserror = "2.0.3"
+rand_chacha = "0.9.0"
+rayon = "1.10.0"
 # Keep arrow in sync with nuts-rs requirements
-arrow = { version = "52.0.0", default-features = false, features = ["ffi"] }
+arrow = { version = "54.1.0", default-features = false, features = ["ffi"] }
 anyhow = "1.0.72"
 itertools = "0.14.0"
 bridgestan = "2.6.1"
-rand_distr = "0.4.3"
-smallvec = "1.11.0"
-upon = { version = "0.8.1", default-features = false, features = [] }
+rand_distr = "0.5.0"
+smallvec = "1.13.0"
+upon = { version = "0.9.0", default-features = false, features = [] }
 time-humanize = { version = "0.1.3", default-features = false }
 indicatif = "0.17.8"
+tch = { version = "0.19.0", optional = true }
 
 [dependencies.pyo3]
-version = "0.21.0"
+version = "0.23.4"
 features = ["extension-module", "anyhow"]
 
 [dev-dependencies]
diff --git a/book.toml b/book.toml
new file mode 100644
index 0000000..9835ea7
--- /dev/null
+++ b/book.toml
@@ -0,0 +1,6 @@
+[book]
+authors = ["Adrian Seyboldt"]
+language = "en"
+multilingual = false
+src = "docs"
+title = "nutpie"
diff --git a/docs/_quarto.yml b/docs/_quarto.yml
new file mode 100644
index 0000000..d443541
--- /dev/null
+++ b/docs/_quarto.yml
@@ -0,0 +1,28 @@
+project:
+  type: website
+
+website:
+  title: "Nutpie"
+  navbar:
+    left:
+      - href: index.qmd
+        text: Home
+      - href: pymc-usage.qmd
+        text: Usage with PyMC
+      - href: stan-usage.qmd
+        text: Usage with Stan
+      - href: sampling-options.qmd
+        text: Sampling Options
+      - href: nf-adapt.qmd
+        text: Normalizing flow adaptation
+      - href: sample-stats.qmd
+        text: Diagnostic information
+      - about.qmd
+
+format:
+  html:
+    theme:
+      - cosmo
+      - brand
+    css: styles.css
+    toc: true
diff --git a/docs/about.qmd b/docs/about.qmd
new file mode 100644
index 0000000..07c5e7f
--- /dev/null
+++ b/docs/about.qmd
@@ -0,0 +1,5 @@
+---
+title: "About"
+---
+
+About this site
diff --git a/docs/benchmarks.md b/docs/benchmarks.md
new file mode 100644
index 0000000..680d565
--- /dev/null
+++ b/docs/benchmarks.md
@@ -0,0 +1 @@
+# Benchmarks
diff --git a/docs/index.qmd b/docs/index.qmd
new file mode 100644
index 0000000..752d937
--- /dev/null
+++ b/docs/index.qmd
@@ -0,0 +1,84 @@
+# Nutpie Documentation
+
+`nutpie` is a high-performance library designed for Bayesian inference, that
+provides efficient sampling algorithms for probabilistic models. It can sample
+models that are defined in PyMC or Stan (numpyro and custom hand-coded
+likelihoods with gradient are coming soon).
+
+- Faster sampling than either the PyMC or Stan default samplers. (An average
+  ~2x speedup on `posteriordb` compared to Stan)
+- All the diagnostic information of PyMC and Stan and some more.
+- GPU support for PyMC models through jax.
+- A more informative progress bar.
+- Access to the incomplete trace during sampling.
+- *Experimental* normalizing flow adaptation for more efficient sampling of
+  difficult posteriors.
+
+## Quickstart
+
+Install `nutpie` with pip, uv, pixi, or conda:
+
+For usage with pymc:
+
+```bash
+# One of
+pip install "nutpie[pymc]"
+uv add "nutpie[pymc]"
+pixi add nutpie pymc numba
+conda install -c conda-forge nutpie pymc numba
+```
+
+And then sample with
+
+```{python}
+import nutpie
+import pymc as pm
+
+with pm.Model() as model:
+    mu = pm.Normal("mu", mu=0, sigma=1)
+    obs = pm.Normal("obs", mu=mu, sigma=1, observed=[1, 2, 3])
+
+compiled = nutpie.compile_pymc_model(model)
+trace = nutpie.sample(compiled)
+```
+
+Stan needs access to a compiler toolchain, you can find instructions for those
+[here](https://mc-stan.org/docs/cmdstan-guide/installation.html#cpp-toolchain).
+You can then install nutpie through pip or uv:
+
+```bash
+# One of
+pip install "nutpie[stan]"
+uv add "nutpie[stan]"
+```
+
+```{python}
+#| echo: false
+import os
+os.environ["TBB_CXX_TYPE"] = "clang"
+```
+
+```{python}
+import nutpie
+
+model = """
+data {
+    int<lower=0> N;
+    vector[N] y;
+}
+parameters {
+    real mu;
+}
+model {
+    mu ~ normal(0, 1);
+    y ~ normal(mu, 1);
+}
+"""
+
+compiled = (
+    nutpie
+    .compile_stan_model(code=model)
+    .with_data(N=3, y=[1, 2, 3])
+)
+trace = nutpie.sample(compiled)
+```
diff --git a/docs/nf-adapt.qmd b/docs/nf-adapt.qmd
new file mode 100644
index 0000000..f6fa300
--- /dev/null
+++ b/docs/nf-adapt.qmd
@@ -0,0 +1,3 @@
+# Adaptation with normalizing flows
+
+**Experimental and subject to change**
diff --git a/docs/pymc-usage.qmd b/docs/pymc-usage.qmd
new file mode 100644
index 0000000..e623eb4
--- /dev/null
+++ b/docs/pymc-usage.qmd
@@ -0,0 +1,177 @@
+# Usage with PyMC models
+
+This document shows how to use `nutpie` with PyMC models. We will use the
+`pymc` package to define a simple model and sample from it using `nutpie`.
+
+## Installation
+
+The recommended way to install `pymc` is through the `conda` ecosystem. A good
+package manager for conda packages is `pixi`. See for the [pixi
+documentation](https://pixi.sh) for instructions on how to install it.
+
+We create a new project for this example:
+
+```bash
+pixi new pymc-example
+```
+
+This will create a new directory `pymc-example` with a `pixi.toml` file, that
+you can edit to add meta information.
+
+We then add the `pymc` and `nutpie` packages to the project:
+
+```bash
+cd pymc-example
+pixi add pymc nutpie arviz
+```
+
+You can use Visual Studio Code (VSCode) or JupyterLab to write and run our code.
+Both are excellent tools for working with Python and data science projects.
+
+### Using VSCode
+
+1. Open VSCode.
+2. Open the `pymc-example` directory created earlier.
+3. Create a new file named `model.ipynb`.
+4. Select the pixi kernel to run the code.
+
+### Using JupyterLab
+
+1. Add jupyter labs to the project by running `pixi add jupyterlab`.
+1. Open JupyterLab by running `pixi run jupyter lab` in your terminal.
+3. Create a new Python notebook.
+
+## Defining and Sampling a Simple Model
+
+We will define a simple Bayesian model using `pymc` and sample from it using
+`nutpie`.
+
+### Model Definition
+
+In your `model.ipypy` file or Jupyter notebook, add the following code:
+
+```{python}
+import pymc as pm
+import nutpie
+import pandas as pd
+
+coords = {"observation": range(3)}
+
+with pm.Model(coords=coords) as model:
+    # Prior distributions for the intercept and slope
+    intercept = pm.Normal("intercept", mu=0, sigma=1)
+    slope = pm.Normal("slope", mu=0, sigma=1)
+
+    # Likelihood (sampling distribution) of observations
+    x = [1, 2, 3]
+
+    mu = intercept + slope * x
+    y = pm.Normal("y", mu=mu, sigma=0.1, observed=[1, 2, 3], dims="observation")
+```
+
+### Sampling
+
+We can now compile the model using the numba backend:
+
+```{python}
+compiled = nutpie.compile_pymc_model(model)
+trace = nutpie.sample(compiled)
+```
+
+While sampling, nutpie shows a progress bar for each chain. It also includes
+information about how each chain is doing:
+
+- It shows the current number of draws
+- The step size of the integrator (very small stepsizes are typically a bad
+  sign)
+- The number of divergences (if there are divergences, that means that nutpie is
+  probably not sampling the posterior correctly)
+- The number of gradient evaluation nutpie uses for each draw. Large numbers
+  (100 to 1000) are a sign that the parameterization of the model is not ideal,
+  and the sampler is very inefficient.
+
+After sampling, this returns an `arviz` InferenceData object that you can use to
+analyze the trace.
+
+For example, we should check the effective sample size:
+
+```{python}
+import arviz as az
+az.ess(trace)
+```
+
+and take a look at a trace plot:
+
+```{python}
+az.plot_trace(trace);
+```
+
+### Choosing the backend
+
+Right now, we have been using the numba backend. This is the default backend for
+`nutpie`, when sampling from pymc models. It tends to have relatively long
+compilation times, but samples small models very efficiently. For larger models
+the `jax` backend sometimes outperforms `numba`.
+
+First, we need to install the `jax` package:
+
+```bash
+pixi add jax
+```
+
+We can select the backend by passing the `backend` argument to the `compile_pymc_model`:
+
+```python
+compiled_jax = nutpie.compiled_pymc_model(model, backend="jax")
+trace = nutpie.sample(compiled_jax)
+```
+
+If you have an nvidia GPU, you can also use the `jax` backend with the `gpu`. We
+will have to install the `jaxlib` package with the `cuda` option
+
+```bash
+pixi add jaxlib --build 'cuda12'
+```
+
+Restart the kernel and check that the GPU is available:
+
+```python
+import jax
+
+# Should list the cuda device
+jax.devices()
+```
+
+Sampling again, should now use the GPU, which you can observe by checking the
+GPU usage with `nvidia-smi` or `nvtop`.
+
+### Changing the dataset without recompilation
+
+If you want to use the same model with different datasets, you can modify
+datasets after compilation. Since jax does not like changes in shapes, this is
+only recommended with the numba backend.
+
+First, we define the model, but put our dataset in a `pm.Data` structure:
+
+```{python}
+with pm.Model() as model:
+    x = pm.Data("x", [1, 2, 3])
+    intercept = pm.Normal("intercept", mu=0, sigma=1)
+    slope = pm.Normal("slope", mu=0, sigma=1)
+    mu = intercept + slope * x
+    y = pm.Normal("y", mu=mu, sigma=0.1, observed=[1, 2, 3])
+```
+
+We can now compile the model:
+
+```{python}
+compiled = nutpie.compile_pymc_model(model)
+trace = nutpie.sample(compiled)
+```
+
+After compilation, we can change the dataset:
+
+```{python}
+compiled2 = compiled.with_data(x=[4, 5, 6])
+trace2 = nutpie.sample(compiled2)
+```
diff --git a/docs/sample-stats.qmd b/docs/sample-stats.qmd
new file mode 100644
index 0000000..e69de29
diff --git a/docs/stan-usage.qmd b/docs/stan-usage.qmd
new file mode 100644
index 0000000..7296231
--- /dev/null
+++ b/docs/stan-usage.qmd
@@ -0,0 +1,230 @@
+# Usage with Stan models
+
+This document shows how to use `nutpie` with Stan models. We will use the
+`nutpie` package to define a simple model and sample from it using Stan.
+
+## Installation
+
+For Stan, it is more common to use `pip` or `uv` to install the necessary
+packages. However, `conda` is also an option if you prefer.
+
+To install using `pip`:
+
+```bash
+pip install "nutpie[stan]"
+```
+
+To install using `uv`:
+
+```bash
+uv add "nutpie[stan]"
+```
+
+To install using `conda`:
+
+```bash
+conda install -c conda-forge nutpie
+```
+
+## Compiler Toolchain
+
+Stan requires a compiler toolchain to be installed on your system. This is
+necessary for compiling the Stan models. You can find detailed instructions for
+setting up the compiler toolchain in the [CmdStan
+Guide](https://mc-stan.org/docs/cmdstan-guide/installation.html#cpp-toolchain).
+
+Additionally, since Stan uses Intel's Threading Building Blocks (TBB) for
+parallelism, you might need to set the `TBB_CXX_TYPE` environment variable to
+specify the compiler type. Depending on your system, you can set it to either
+`clang` or `gcc`. For example:
+
+```{python}
+import os
+os.environ["TBB_CXX_TYPE"] = "clang"  # or 'gcc'
+```
+
+Make sure to set this environment variable before compiling your Stan models to ensure proper configuration.
+
+## Defining and Sampling a Simple Model
+
+We will define a simple Bayesian model using Stan and sample from it using
+`nutpie`.
+
+### Model Definition
+
+In your Python script or Jupyter notebook, add the following code:
+
+```{python}
+import nutpie
+
+model_code = """
+data {
+    int<lower=0> N;
+    vector[N] y;
+}
+parameters {
+    real mu;
+}
+model {
+    mu ~ normal(0, 1);
+    y ~ normal(mu, 1);
+}
+"""
+
+compiled_model = nutpie.compile_stan_model(code=model_code)
+```
+
+### Sampling
+
+We can now compile the model and sample from it:
+
+```{python}
+compiled_model_with_data = compiled_model.with_data(N=3, y=[1, 2, 3])
+trace = nutpie.sample(compiled_model_with_data)
+```
+
+### Using Dimensions
+
+We'll use the radon model from
+[this](https://mc-stan.org/learn-stan/case-studies/radon_cmdstanpy_plotnine.html)
+case-study from the stan documentation, to show how we can use coordinates and
+dimension names to simplify working with trace objects.
+
+We follow the same data preparation as in the case-study:
+
+```{python}
+import pandas as pd
+import numpy as np
+import arviz as az
+import seaborn as sns
+
+home_data = pd.read_csv(
+    "https://github.com/pymc-devs/pymc-examples/raw/refs/heads/main/examples/data/srrs2.dat",
+    index_col="idnum",
+)
+county_data = pd.read_csv(
+    "https://github.com/pymc-devs/pymc-examples/raw/refs/heads/main/examples/data/cty.dat",
+)
+
+radon_data = (
+    home_data
+    .rename(columns=dict(cntyfips="ctfips"))
+    .merge(
+        (
+            county_data
+            .drop_duplicates(['stfips', 'ctfips', 'st', 'cty', 'Uppm'])
+            .set_index(["ctfips", "stfips"])
+        ),
+        right_index=True,
+        left_on=["ctfips", "stfips"],
+    )
+    .assign(log_radon=lambda x: np.log(np.clip(x.activity, 0.1, np.inf)))
+    .assign(log_uranium=lambda x: np.log(np.clip(x["Uppm"], 0.1, np.inf)))
+    .query("state == 'MN'")
+)
+```
+
+And also use the partially pooled model from the case-study:
+
+```{python}
+model_code = """
+data {
+  int<lower=1> N;  // observations
+  int<lower=1> J;  // counties
+  array[N] int<lower=1, upper=J> county;
+  vector[N] x;
+  vector[N] y;
+}
+parameters {
+  real mu_alpha;
+  real<lower=0> sigma_alpha;
+  vector<offset=mu_alpha, multiplier=sigma_alpha>[J] alpha;  // non-centered parameterization
+  real beta;
+  real<lower=0> sigma;
+}
+model {
+  y ~ normal(alpha[county] + beta * x, sigma);
+  alpha ~ normal(mu_alpha, sigma_alpha); // partial-pooling
+  beta ~ normal(0, 10);
+  sigma ~ normal(0, 10);
+  mu_alpha ~ normal(0, 10);
+  sigma_alpha ~ normal(0, 10);
+}
+generated quantities {
+  array[N] real y_rep = normal_rng(alpha[county] + beta * x, sigma);
+}
+"""
+```
+
+We collect the dataset in the format that the stan model requires,
+and specify the dimensions of each of the non-scalar variables in the model:
+
+```{python}
+county_idx, counties = pd.factorize(radon_data["county"], use_na_sentinel=False)
+observations = radon_data.index
+
+coords = {
+    "county": counties,
+    "observation": observations,
+}
+
+dims = {
+    "alpha": ["county"],
+    "y_rep": ["observation"],
+}
+
+data = {
+    "N": len(observations),
+    "J": len(counties),
+    # Stan uses 1-based indexing!
+    "county": county_idx + 1,
+    "x": radon_data.log_uranium.values,
+    "y": radon_data.log_radon.values,
+}
+```
+
+Then, we compile the model and provide the dimensions, coordinates and the
+dataset we just defined:
+
+```{python}
+compiled_model = (
+    nutpie.compile_stan_model(code=model_code)
+    .with_data(**data)
+    .with_dims(**dims)
+    .with_coords(**coords)
+)
+```
+
+```{python}
+%%time
+trace = nutpie.sample(compiled_model, seed=0)
+```
+
+As some basic convergance checking we verify that all Rhat values are smaller
+than 1.02, all parameters have at least 500 effective draws and that we have no
+divergences:
+
+```{python}
+assert trace.sample_stats.diverging.sum() == 0
+assert az.ess(trace).min().min() > 500
+assert az.rhat(trace).max().max() > 1.02
+```
+
+Thanks to the coordinates and dimensions we specified, the resulting trace will
+now contain labeled data, so that plots based on it have properly set-up labels:
+
+```{python}
+import arviz as az
+import seaborn as sns
+import xarray as xr
+
+sns.catplot(
+    data=trace.posterior.alpha.to_dataframe().reset_index(),
+    y="county",
+    x="alpha",
+    kind="boxen",
+    height=13,
+    aspect=1/2.5,
+    showfliers=False,
+)
+```
diff --git a/docs/styles.css b/docs/styles.css
new file mode 100644
index 0000000..2ddf50c
--- /dev/null
+++ b/docs/styles.css
@@ -0,0 +1 @@
+/* css styles */
diff --git a/pixi.lock b/pixi.lock
new file mode 100644
index 0000000..da15dad
--- /dev/null
+++ b/pixi.lock
@@ -0,0 +1,4474 @@
+version: 6
+environments:
+  default:
+    channels:
+    - url: https://conda.anaconda.org/conda-forge/
+    indexes:
+    - https://pypi.org/simple
+    packages:
+      linux-64:
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_kmp_llvm.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/adwaita-icon-theme-47.0-unix_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/arviz-0.20.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/asciitree-0.3.3-py_2.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/at-spi2-atk-2.38.0-h0630a04_3.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/at-spi2-core-2.40.3-h0630a04_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/atk-1.0-2.38.0-h04ea711_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.1.0-pyh71513ae_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.43-h4bf12b8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_linux-64-2.43-h4852527_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/blas-2.130-blis.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/blas-devel-3.9.0-30_hdec4247_blis.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/blis-0.9.0-h4ab18f5_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-h4bc722e_7.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.4-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2025.1.31-hbcca054_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cachetools-5.5.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cairo-1.18.2-h3394656_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cloudpickle-3.1.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cmdstan-2.36.0-h6b7b5dd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cmdstanpy-1.2.5-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.2-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cons-0.4.6-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.1-py312h68727a3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/dart-sass-1.58.3-ha770c72_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/dbus-1.13.6-h5008d03_3.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.12-py312h2ec8cdc_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.1.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/deno-1.46.3-hcab8b69_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/deno-dom-0.1.41-h4768de7_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/deprecated-1.2.18-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/epoxy-1.5.10-h166bdaf_1.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/esbuild-0.24.2-ha770c72_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/etuples-0.3.9-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/exceptiongroup-1.2.2-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.1.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/expat-2.6.4-h5888daf_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fasteners-0.19-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.17.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-dejavu-sans-mono-2.37-hab24e00_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-inconsolata-3.000-h77eed37_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-source-code-pro-2.038-h77eed37_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-ubuntu-0.83-h77eed37_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/fontconfig-2.15.0-h7e30c49_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-ecosystem-1-0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-forge-1-0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.56.0-py312h178313f_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.12.1-h267a509_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/fribidi-1.0.10-h36c2ea0_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc-13.3.0-h9576a4e_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_impl_linux-64-13.3.0-hfea6d02_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_linux-64-13.3.0-hc28eda2_7.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gdk-pixbuf-2.42.12-hb9ae30d_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/glib-tools-2.82.2-h4833e2c_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/graphite2-1.3.13-h59595ed_1003.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/graphviz-12.2.1-h5ae0cbf_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gtk3-3.24.43-h021d004_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gts-0.7.6-h977cf35_4.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gxx-13.3.0-h9576a4e_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gxx_impl_linux-64-13.3.0-hdbfa832_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gxx_linux-64-13.3.0-h6834431_7.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/h5netcdf-1.5.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/h5py-3.13.0-nompi_py312hedeef09_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/harfbuzz-10.2.0-h4bba637_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/hdf5-1.14.3-nompi_h2d575fe_109.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/hicolor-icon-theme-0.17-ha770c72_2.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-75.1-he02047a_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/importlib_resources-6.5.2-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/ipykernel-6.29.5-pyh3099207_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/ipython-8.32.0-pyh907856f_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/ipywidgets-8.1.5-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jsonschema-4.23.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jsonschema-specifications-2024.10.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_client-8.6.3-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_core-5.7.2-pyh31011fe_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_widgets-3.0.13-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/kernel-headers_linux-64-3.10.0-he073ed8_18.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.1-h166bdaf_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.8-py312h84d6215_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.43-h712a8e2_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h27087fc_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libaec-1.1.3-h59595ed_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-30_h66dfbfd_blis.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-30_hba4ea11_blis.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libcups-2.3.3-h4637d8d_4.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.12.1-h332b0f4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.23-h4ddbbb0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h7949ede_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-hd590300_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.6.4-h5888daf_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.4.6-h2dba641_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-14.2.0-h77fa898_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/libgcc-devel_linux-64-13.3.0-h84ea5a7_101.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-14.2.0-h69a702a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgd-2.3.3-h6f5c62b_11.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-14.2.0-h69a702a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-14.2.0-hd5240d6_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libglib-2.82.2-h2ff4ddf_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-14.2.0-h77fa898_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libhwloc-2.11.2-default_h0d58e46_1001.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h4ce23a2_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.0.0-hd590300_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-9_h9f1adc1_netlib.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapacke-3.9.0-9_h0ad7b2f_netlib.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libllvm14-14.0.6-hcd5def8_4.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.6.4-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.64.0-h161d5f1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hd590300_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.47-h943b412_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/librsvg-2.58.4-h49af25d_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libsanitizer-13.3.0-heb74ff8_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libsodium-1.0.20-h4ab18f5_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.49.1-hee588c1_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-hf672d98_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-14.2.0-hc0a3c3a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/libstdcxx-devel_linux-64-13.3.0-h84ea5a7_101.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-14.2.0-h4852527_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.0-hd9ff511_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.38.1-h0b41bf4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.5.0-h851e524_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libxkbcommon-1.8.0-hc4a0caf_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.13.6-h8d12d68_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/llvm-openmp-19.1.7-h024ca30_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/llvmlite-0.43.0-py312h374181b_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/logical-unification-0.4.6-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/make-4.4.1-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-3.0.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.0-py312hd3ec401_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.1.7-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/maturin-1.8.2-py312h6ab59e4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/minikanren-1.0.3-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/mkl-2024.2.2-ha957f24_16.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/mkl-service-2.4.2-py312hf224ee7_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/msgpack-python-1.1.0-py312h68727a3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/multipledispatch-0.6.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyh9f0ad1d_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/nbclient-0.10.2-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/nbformat-5.10.4-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/nest-asyncio-1.6.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numba-0.60.0-py312h83e6fd3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numcodecs-0.15.1-py312hf9745cd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py312heda63a1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/objprint-0.3.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.3-h5fbd93e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.4.1-h7b32b05_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.10.15-py312h12e396e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-24.2-pyhd8ed1ab_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-2.2.3-py312hf9745cd_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pandoc-3.4-ha770c72_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pango-1.56.1-h861ebed_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.4-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/patsy-1.0.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.44-hba22ea6_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pickleshare-0.7.5-pyhd8ed1ab_1004.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-11.1.0-py312h80c1187_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pip-24.3.1-pyh8b19718_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pixman-0.44.2-h29eaf8c_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pkgutil-resolve-name-1.3.10-pyhd8ed1ab_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.3.6-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.5.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/polars-1.22.0-py312hda0fa55_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.50-pyha770c72_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-6.1.1-py312h66e93f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.19.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pymc-5.20.1-hd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pymc-base-5.20.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.2.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pytensor-2.27.1-py312h97902ae_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pytensor-base-2.27.1-py312h25a0e75_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-8.3.4-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.9-h9e4cc4f_0_cpython.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhff2d567_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-graphviz-0.20.3-pyh91182bf_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python_abi-3.12-5_cp312.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2024.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.2-py312h178313f_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-26.2.1-py312hbf22597_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/quarto-1.6.40-ha770c72_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.2-h8228510_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/referencing-0.36.2-pyh29332c3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/rich-13.9.4-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.22.3-py312h12e396e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.15.2-py312ha707e6e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-0.13.2-hd8ed1ab_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-base-0.13.2-pyhd8ed1ab_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-75.8.0-pyhff2d567_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/stanio-0.5.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/statsmodels-0.14.4-py312hc0a28a1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/sysroot_linux-64-2.17-h0157908_18.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tbb-2021.13.0-hceb3a55_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tbb-devel-2021.13.0-h1f99690_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/threadpoolctl-3.5.0-pyhc1e730c_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h4845f30_101.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.2.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/toolz-1.0.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.4.2-py312h66e93f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/tqdm-4.67.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.12.2-hd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.12.2-pyha770c72_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/typst-0.11.0-he8a937b_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025a-h78e105d_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-16.0.0-py312h66e93f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/viztracer-1.0.2-py312h0f154a2_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/wayland-1.23.1-h3e06ad9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.13-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/wheel-0.45.1-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/widgetsnbextension-4.0.13-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/wrapt-1.17.2-py312h66e93f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/xarray-2025.1.2-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/xarray-einstats-0.8.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xkeyboard-config-2.43-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libice-1.1.2-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libsm-1.2.5-he73a12e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libx11-1.8.11-h4f16b4b_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcomposite-0.4.6-hb9d3cd8_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcursor-1.2.3-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdamage-1.1.6-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxext-1.3.6-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxfixes-6.0.1-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxi-1.8.2-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxinerama-1.1.5-h5888daf_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrandr-1.5.4-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrender-0.9.12-hb9d3cd8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxtst-1.2.5-hb9d3cd8_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h7f98852_2.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/zarr-2.18.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/zeromq-4.3.5-h3b0a872_7.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.21.0-pyhd8ed1ab_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb8e6e7a_0.conda
+      - pypi: https://files.pythonhosted.org/packages/a2/ad/e0d3c824784ff121c03cc031f944bc7e139a8f1870ffd2845cc2dd76f6c4/absl_py-2.1.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/8c/ce/b9a154415fe7f51a481bece2bd998c679ce0ef93833ad4e7c87c60f0e461/bridgestan-2.6.1-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/38/fc/bce832fd4fd99766c04d1ee0eead6b0ec6486fb100ae5e74c1d91292b982/certifi-2025.1.31-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/3e/a2/513f6cbe752421f16d969e32f3583762bfd583848b763913ddab8d9bfd4f/charset_normalizer-3.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/58/a0/e4bebe76bdd0a68077030f1b5e48b545597473ae1b773c84150311152efc/chex-0.1.88-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/7b/07/a212ca4cfe56fa35c8315307e37df218e5946c726a792b5c7a795b245c10/dllist-2.0.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/2c/90dd8d90acae15cab7ad344398030a0ba6a1cb1257d5e7da57f41e566940/equinox-0.11.11-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/05/83/bb4a4518bfa32a160dc455d8d944a6e00a7eb6759f449cb20c7b7879090e/etils-1.12.0-py3-none-any.whl
+      - pypi: git+https://github.com/aseyboldt/flowjax?rev=07e7e32217bcfcaa7d68a304f332c925d26ab76f#07e7e32217bcfcaa7d68a304f332c925d26ab76f
+      - pypi: https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/f4/58/cc0721a1030fcbab0984beea0bf3c4610ec103f738423cdfa9c4ceb40598/jax-0.5.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/66/e9/211ba3e46ec22c722c4d61a739cfccf79b0618006d6f5fa53eb4eb93ed6d/jaxlib-0.5.0-cp312-cp312-manylinux2014_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/db/7e/da7b57a1f3af7303a0f3c8594d820fc0d3a9bbe3810a357eb21eb166e76b/jaxtyping-0.2.38-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/08/57/5d58fad4124192b1be42f68bd0c0ddaa26e44a730ff8c9337adade2f5632/ml_dtypes-0.5.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/5c/24/28d0bb21600a78e46754947333ec9a297044af884d360092eb8561575fe9/optax-0.2.4-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/f8/7b1f5b84f84e43face7349ac2df23f0f75d36b5b4a1c1c0305edce82bcc8/paramax-0.0.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/8f/4d/b72e0782abec07f3d8dabf24cf12673d26b173af2046eb4e67365c776ccf/posteriordb-0.2.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/f9/9b/335f9764261e915ed497fcdeb11df5dfd6f7bf257d4a6a2a686d80da4d54/requests-2.32.3-py3-none-any.whl
+      - pypi: git+https://github.com/aseyboldt/samplerlab#9edea2da1c5b336e53b9238440a2317cf1baba14
+      - pypi: https://files.pythonhosted.org/packages/c8/19/4ec628951a74043532ca2cf5d97b7b14863931476d117c471e8e2b1eb39f/urllib3-2.3.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/39/3b/5b918a0da0d6920e7f7328cf0ab00df31b905d709f458596304f09096785/wadler_lindig-0.1.3-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/c3/ac/1847959031ce09eb65a70b16d0f07633d7ddaf8a97752cc924f9a8cbe253/watermark-2.5.0-py2.py3-none-any.whl
+packages:
+- conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
+  sha256: fe51de6107f9edc7aa4f786a70f4a883943bc9d39b3bb7307c04c41410990726
+  md5: d7c89558ba9fa0495403155b64376d81
+  arch: x86_64
+  platform: linux
+  license: None
+  purls: []
+  size: 2562
+  timestamp: 1578324546067
+- conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_kmp_llvm.tar.bz2
+  build_number: 2
+  sha256: 84a66275da3a66e3f3e70e9d8f10496d807d01a9e4ec16cd2274cc5e28c478fc
+  md5: 562b26ba2e19059551a811e72ab7f793
+  depends:
+  - _libgcc_mutex 0.1 conda_forge
+  - llvm-openmp >=9.0.1
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 5744
+  timestamp: 1650742457817
+- pypi: https://files.pythonhosted.org/packages/a2/ad/e0d3c824784ff121c03cc031f944bc7e139a8f1870ffd2845cc2dd76f6c4/absl_py-2.1.0-py3-none-any.whl
+  name: absl-py
+  version: 2.1.0
+  sha256: 526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308
+  requires_python: '>=3.7'
+- conda: https://conda.anaconda.org/conda-forge/noarch/adwaita-icon-theme-47.0-unix_0.conda
+  sha256: 188dca9a847f474b3df71eda9fe828fbe10b53aa6f4313c7e117f3114b1dd84e
+  md5: 49436a5c604f99058473d84580f0e341
+  depends:
+  - __unix
+  - hicolor-icon-theme
+  - librsvg
+  license: LGPL-3.0-or-later OR CC-BY-SA-3.0
+  license_family: LGPL
+  purls: []
+  size: 566980
+  timestamp: 1728314504182
+- conda: https://conda.anaconda.org/conda-forge/noarch/arviz-0.20.0-pyhd8ed1ab_1.conda
+  sha256: a41ee718207a4defccf6eae8e85f622e82fa45323744dd49af8657cb5872ab43
+  md5: 03111f1a24b520153bf7912a36eb65f3
+  depends:
+  - h5netcdf >=1.0.2
+  - matplotlib-base >=3.5
+  - numpy >=1.23.0
+  - packaging
+  - pandas >=1.5.0
+  - python >=3.10
+  - scipy >=1.9.0
+  - setuptools >=60.0.0
+  - typing_extensions >=4.1.0
+  - xarray >=2022.6.0
+  - xarray-einstats >=0.3
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/arviz?source=hash-mapping
+  size: 1473775
+  timestamp: 1734087589897
+- conda: https://conda.anaconda.org/conda-forge/noarch/asciitree-0.3.3-py_2.tar.bz2
+  sha256: b3e9369529fe7d721b66f18680ff4b561e20dbf6507e209e1f60eac277c97560
+  md5: c0481c9de49f040272556e2cedf42816
+  depends:
+  - python
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/asciitree?source=hash-mapping
+  size: 6164
+  timestamp: 1531050741142
+- conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.0-pyhd8ed1ab_1.conda
+  sha256: 93b14414b3b3ed91e286e1cbe4e7a60c4e1b1c730b0814d1e452a8ac4b9af593
+  md5: 8f587de4bcf981e26228f268df374a9b
+  depends:
+  - python >=3.9
+  constrains:
+  - astroid >=2,<4
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/asttokens?source=hash-mapping
+  size: 28206
+  timestamp: 1733250564754
+- conda: https://conda.anaconda.org/conda-forge/linux-64/at-spi2-atk-2.38.0-h0630a04_3.tar.bz2
+  sha256: 26ab9386e80bf196e51ebe005da77d57decf6d989b4f34d96130560bc133479c
+  md5: 6b889f174df1e0f816276ae69281af4d
+  depends:
+  - at-spi2-core >=2.40.0,<2.41.0a0
+  - atk-1.0 >=2.36.0
+  - dbus >=1.13.6,<2.0a0
+  - libgcc-ng >=9.3.0
+  - libglib >=2.68.1,<3.0a0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  license_family: LGPL
+  purls: []
+  size: 339899
+  timestamp: 1619122953439
+- conda: https://conda.anaconda.org/conda-forge/linux-64/at-spi2-core-2.40.3-h0630a04_0.tar.bz2
+  sha256: c4f9b66bd94c40d8f1ce1fad2d8b46534bdefda0c86e3337b28f6c25779f258d
+  md5: 8cb2fc4cd6cc63f1369cfa318f581cc3
+  depends:
+  - dbus >=1.13.6,<2.0a0
+  - libgcc-ng >=9.3.0
+  - libglib >=2.68.3,<3.0a0
+  - xorg-libx11
+  - xorg-libxi
+  - xorg-libxtst
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  license_family: LGPL
+  purls: []
+  size: 658390
+  timestamp: 1625848454791
+- conda: https://conda.anaconda.org/conda-forge/linux-64/atk-1.0-2.38.0-h04ea711_2.conda
+  sha256: df682395d05050cd1222740a42a551281210726a67447e5258968dd55854302e
+  md5: f730d54ba9cd543666d7220c9f7ed563
+  depends:
+  - libgcc-ng >=12
+  - libglib >=2.80.0,<3.0a0
+  - libstdcxx-ng >=12
+  constrains:
+  - atk-1.0 2.38.0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.0-or-later
+  license_family: LGPL
+  purls: []
+  size: 355900
+  timestamp: 1713896169874
+- conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.1.0-pyh71513ae_0.conda
+  sha256: 1f267886522dfb9ae4e5ebbc3135b5eb13cff27bdbfe8d881a4d893459166ab4
+  md5: 2cc3f588512f04f3a0c64b4e9bedc02d
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/attrs?source=compressed-mapping
+  size: 56370
+  timestamp: 1737819298139
+- conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.43-h4bf12b8_2.conda
+  sha256: 267e78990247369b13234bda270f31beb56a600b4851a8244e31dd9ad85b3b17
+  md5: cf0c5521ac2a20dfa6c662a4009eeef6
+  depends:
+  - ld_impl_linux-64 2.43 h712a8e2_2
+  - sysroot_linux-64
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only
+  license_family: GPL
+  purls: []
+  size: 5682777
+  timestamp: 1729655371045
+- conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_linux-64-2.43-h4852527_2.conda
+  sha256: df52bd8b8b2a20a0c529d9ad08aaf66093ac318aa8a33d270f18274341a77062
+  md5: 18aba879ddf1f8f28145ca6fcb873d8c
+  depends:
+  - binutils_impl_linux-64 2.43 h4bf12b8_2
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only
+  license_family: GPL
+  purls: []
+  size: 34945
+  timestamp: 1729655404893
+- conda: https://conda.anaconda.org/conda-forge/linux-64/blas-2.130-blis.conda
+  build_number: 30
+  sha256: 2b729e6a31be5a82e2885b05e9ac7a1bbc6dd99f368685bd71d581ae0052f4df
+  md5: 9dafabe6baf2f6660a0922311e97aabd
+  depends:
+  - blas-devel 3.9.0 30*_blis
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 16759
+  timestamp: 1739836358004
+- conda: https://conda.anaconda.org/conda-forge/linux-64/blas-devel-3.9.0-30_hdec4247_blis.conda
+  build_number: 30
+  sha256: 320cfbe67fac6dc33445386372f07029706bd9422765d4dee2e006034de9843c
+  md5: 75470f91acc0b1f1966837d5b6ff76f5
+  depends:
+  - blis 0.9.0.*
+  - libblas 3.9.0 30_h66dfbfd_blis
+  - libcblas 3.9.0 30_hba4ea11_blis
+  - liblapack 3.9.0 *_netlib
+  - liblapacke 3.9.0 *_netlib
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 16561
+  timestamp: 1739836226671
+- conda: https://conda.anaconda.org/conda-forge/linux-64/blis-0.9.0-h4ab18f5_2.conda
+  sha256: 3e501cbf98ccb69210e6145d38295dc14ca11417e9c86fec988f06adea8456fd
+  md5: 6f77ba1352b69c4a6f8a6d20def30e4e
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 3547115
+  timestamp: 1713877874618
+- pypi: https://files.pythonhosted.org/packages/8c/ce/b9a154415fe7f51a481bece2bd998c679ce0ef93833ad4e7c87c60f0e461/bridgestan-2.6.1-py3-none-any.whl
+  name: bridgestan
+  version: 2.6.1
+  sha256: 142ea79542db6022e8d85b8d5eaed7817f0429ad0156dbef7e147c72e72bcd98
+  requires_dist:
+  - numpy
+  - dllist
+  - pandas ; extra == 'test'
+  - pytest ; extra == 'test'
+  - pytest-cov ; extra == 'test'
+  requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hb9d3cd8_2.conda
+  sha256: fcb0b5b28ba7492093e54f3184435144e074dfceab27ac8e6a9457e736565b0b
+  md5: 98514fe74548d768907ce7a13f680e8f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - brotli-bin 1.1.0 hb9d3cd8_2
+  - libbrotlidec 1.1.0 hb9d3cd8_2
+  - libbrotlienc 1.1.0 hb9d3cd8_2
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 19264
+  timestamp: 1725267697072
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hb9d3cd8_2.conda
+  sha256: 261364d7445513b9a4debc345650fad13c627029bfc800655a266bf1e375bc65
+  md5: c63b5e52939e795ba8d26e35d767a843
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libbrotlidec 1.1.0 hb9d3cd8_2
+  - libbrotlienc 1.1.0 hb9d3cd8_2
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 18881
+  timestamp: 1725267688731
+- conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-h4bc722e_7.conda
+  sha256: 5ced96500d945fb286c9c838e54fa759aa04a7129c59800f0846b4335cee770d
+  md5: 62ee74e96c5ebb0af99386de58cf9553
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: bzip2-1.0.6
+  license_family: BSD
+  purls: []
+  size: 252783
+  timestamp: 1720974456583
+- conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.4-hb9d3cd8_0.conda
+  sha256: d4f28d87b6339b94f74762c0076e29c8ef8ddfff51a564a92da2843573c18320
+  md5: e2775acf57efd5af15b8e3d1d74d72d3
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 206085
+  timestamp: 1734208189009
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2025.1.31-hbcca054_0.conda
+  sha256: bf832198976d559ab44d6cdb315642655547e26d826e34da67cbee6624cda189
+  md5: 19f3a56f68d2fd06c516076bff482c52
+  arch: x86_64
+  platform: linux
+  license: ISC
+  purls: []
+  size: 158144
+  timestamp: 1738298224464
+- conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
+  noarch: python
+  sha256: 561e6660f26c35d137ee150187d89767c988413c978e1b712d53f27ddf70ea17
+  md5: 9b347a7ec10940d3f7941ff6c460b551
+  depends:
+  - cached_property >=1.5.2,<1.5.3.0a0
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 4134
+  timestamp: 1615209571450
+- conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
+  sha256: 6dbf7a5070cc43d90a1e4c2ec0c541c69d8e30a0e25f50ce9f6e4a432e42c5d7
+  md5: 576d629e47797577ab0f1b351297ef4a
+  depends:
+  - python >=3.6
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/cached-property?source=hash-mapping
+  size: 11065
+  timestamp: 1615209567874
+- conda: https://conda.anaconda.org/conda-forge/noarch/cachetools-5.5.1-pyhd8ed1ab_0.conda
+  sha256: 04cd27394393d5e9c6315e7e6a344ba38ddfa49f899c05643208ccba07968430
+  md5: 6eb7c1074d938746b195f556abf9a28f
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/cachetools?source=hash-mapping
+  size: 14948
+  timestamp: 1737517675303
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cairo-1.18.2-h3394656_1.conda
+  sha256: de7d0d094e53decc005cb13e527be2635b8f604978da497d4c0d282c7dc08385
+  md5: b34c2833a1f56db610aeb27f206d800d
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - fontconfig >=2.15.0,<3.0a0
+  - fonts-conda-ecosystem
+  - freetype >=2.12.1,<3.0a0
+  - icu >=75.1,<76.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libglib >=2.82.2,<3.0a0
+  - libpng >=1.6.44,<1.7.0a0
+  - libstdcxx >=13
+  - libxcb >=1.17.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - pixman >=0.44.2,<1.0a0
+  - xorg-libice >=1.1.1,<2.0a0
+  - xorg-libsm >=1.2.4,<2.0a0
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxrender >=0.9.11,<0.10.0a0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-only or MPL-1.1
+  purls: []
+  size: 978868
+  timestamp: 1733790976384
+- pypi: https://files.pythonhosted.org/packages/38/fc/bce832fd4fd99766c04d1ee0eead6b0ec6486fb100ae5e74c1d91292b982/certifi-2025.1.31-py3-none-any.whl
+  name: certifi
+  version: 2025.1.31
+  sha256: ca78db4565a652026a4db2bcdf68f2fb589ea80d0be70e03929ed730746b84fe
+  requires_python: '>=3.6'
+- pypi: https://files.pythonhosted.org/packages/3e/a2/513f6cbe752421f16d969e32f3583762bfd583848b763913ddab8d9bfd4f/charset_normalizer-3.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+  name: charset-normalizer
+  version: 3.4.1
+  sha256: bc2722592d8998c870fa4e290c2eec2c1569b87fe58618e67d38b4665dfa680d
+  requires_python: '>=3.7'
+- pypi: https://files.pythonhosted.org/packages/58/a0/e4bebe76bdd0a68077030f1b5e48b545597473ae1b773c84150311152efc/chex-0.1.88-py3-none-any.whl
+  name: chex
+  version: 0.1.88
+  sha256: 234b61a5baa8132802e4b9c5657167d6c8a911d90a59a0bec47d537567e41b75
+  requires_dist:
+  - absl-py>=0.9.0
+  - typing-extensions>=4.2.0
+  - jax>=0.4.27
+  - jaxlib>=0.4.27
+  - numpy>=1.24.1
+  - setuptools ; python_full_version >= '3.12'
+  - toolz>=0.9.0
+  requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/cloudpickle-3.1.1-pyhd8ed1ab_0.conda
+  sha256: 21ecead7268241007bf65691610cd7314da68c1f88113092af690203b5780db5
+  md5: 364ba6c9fb03886ac979b482f39ebb92
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/cloudpickle?source=hash-mapping
+  size: 25870
+  timestamp: 1736947650712
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cmdstan-2.36.0-h6b7b5dd_0.conda
+  sha256: d475cfa2ebf3fa4f4d1cda77740f57c890ca851cf7e1b9f511d4f67791b2f0f4
+  md5: 7c833c27c240b57fd29d1543b41f95a8
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - gxx_linux-64 13.*
+  - gxx_linux-64 >=13
+  - libgcc >=13
+  - libstdcxx >=13
+  - make
+  - sysroot_linux-64 2.17.*
+  - tbb >=2021.13.0
+  - tbb-devel
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 32912355
+  timestamp: 1733858441087
+- conda: https://conda.anaconda.org/conda-forge/noarch/cmdstanpy-1.2.5-pyhd8ed1ab_0.conda
+  sha256: 3654832b719b29c7d7c05f7fb031127190ca059575f74310eb55ab9e27e6a87b
+  md5: 4046ad21a0f33a10c1ea4fcbed72856c
+  depends:
+  - cmdstan
+  - numpy >=1.21
+  - pandas
+  - python >=3.9
+  - stanio >=0.4.0,<2.0.0
+  - tqdm
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/cmdstanpy?source=hash-mapping
+  size: 70880
+  timestamp: 1734212972209
+- conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
+  sha256: ab29d57dc70786c1269633ba3dff20288b81664d3ff8d21af995742e2bb03287
+  md5: 962b9857ee8e7018c22f2776ffa0b2d7
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/colorama?source=hash-mapping
+  size: 27011
+  timestamp: 1733218222191
+- conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.2-pyhd8ed1ab_1.conda
+  sha256: 7e87ef7c91574d9fac19faedaaee328a70f718c9b4ddadfdc0ba9ac021bd64af
+  md5: 74673132601ec2b7fc592755605f4c1b
+  depends:
+  - python >=3.9
+  - traitlets >=5.3
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/comm?source=hash-mapping
+  size: 12103
+  timestamp: 1733503053903
+- conda: https://conda.anaconda.org/conda-forge/noarch/cons-0.4.6-pyhd8ed1ab_1.conda
+  sha256: 444f2df4c59f624bf82c9bc23e5d0e4d50f26fbb477197b5b1d2dd151a3bcd69
+  md5: 407ddb4cf0d95f21326af9e3df56627f
+  depends:
+  - logical-unification >=0.4.1
+  - python >=3.9
+  license: LGPL-3.0-only
+  license_family: LGPL
+  purls:
+  - pkg:pypi/cons?source=hash-mapping
+  size: 14424
+  timestamp: 1734526937473
+- conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.1-py312h68727a3_0.conda
+  sha256: e977af50b844b5b8cfec358131a4e923f0aa718e8334321cf8d84f5093576259
+  md5: f5fbba0394ee45e9a64a73c2a994126a
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - numpy >=1.23
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/contourpy?source=hash-mapping
+  size: 276332
+  timestamp: 1731428454756
+- conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhd8ed1ab_1.conda
+  sha256: 9827efa891e507a91a8a2acf64e210d2aff394e1cde432ad08e1f8c66b12293c
+  md5: 44600c4667a319d67dbe0681fc0bc833
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/cycler?source=hash-mapping
+  size: 13399
+  timestamp: 1733332563512
+- conda: https://conda.anaconda.org/conda-forge/linux-64/dart-sass-1.58.3-ha770c72_1.conda
+  sha256: a94f8024ac8a09cbf265c62d2bd57e709802868ef656c6cafe5864ed20bf94af
+  md5: d54982a58cd9be3d00a7efe76ba6f60c
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 2961769
+  timestamp: 1683598364427
+- conda: https://conda.anaconda.org/conda-forge/linux-64/dbus-1.13.6-h5008d03_3.tar.bz2
+  sha256: 8f5f995699a2d9dbdd62c61385bfeeb57c82a681a7c8c5313c395aa0ccab68a5
+  md5: ecfff944ba3960ecb334b9a2663d708d
+  depends:
+  - expat >=2.4.2,<3.0a0
+  - libgcc-ng >=9.4.0
+  - libglib >=2.70.2,<3.0a0
+  arch: x86_64
+  platform: linux
+  license: GPL-2.0-or-later
+  license_family: GPL
+  purls: []
+  size: 618596
+  timestamp: 1640112124844
+- conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.12-py312h2ec8cdc_0.conda
+  sha256: f88c3a7ff384d1726aea2cb2342cf67f1502915391860335c40ab81d7e381e30
+  md5: 6be6dcb4bffd1d456bdad28341d507bd
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/debugpy?source=hash-mapping
+  size: 2646757
+  timestamp: 1737269937348
+- conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.1.1-pyhd8ed1ab_1.conda
+  sha256: 84e5120c97502a3785e8c3241c3bf51f64b4d445f13b4d2445db00d9816fe479
+  md5: d622d8d7ee8868870f9cbe259f381181
+  depends:
+  - python >=3.9
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/decorator?source=hash-mapping
+  size: 14068
+  timestamp: 1733236549190
+- conda: https://conda.anaconda.org/conda-forge/linux-64/deno-1.46.3-hcab8b69_0.conda
+  sha256: 704bb88a744aacab35c159929d6aee8216e7731e920f7d28254a721059dde4b1
+  md5: f8410ffda105ab8d380990bbdf9f266f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 39585861
+  timestamp: 1726025052074
+- conda: https://conda.anaconda.org/conda-forge/linux-64/deno-dom-0.1.41-h4768de7_0.conda
+  sha256: 2726829e1a6117dc9555a2088a4ebaa05035218b8791106f036ba73f15b1b13e
+  md5: 89e80615c6b1b1ae4f5267bfc5fc7f66
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - deno >=1.24.2
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 433150
+  timestamp: 1736703310469
+- conda: https://conda.anaconda.org/conda-forge/noarch/deprecated-1.2.18-pyhd8ed1ab_0.conda
+  sha256: d614bcff10696f1efc714df07651b50bf3808401fcc03814309ecec242cc8870
+  md5: 0cef44b1754ae4d6924ac0eef6b9fdbe
+  depends:
+  - python >=3.9
+  - wrapt <2,>=1.10
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/deprecated?source=hash-mapping
+  size: 14382
+  timestamp: 1737987072859
+- pypi: https://files.pythonhosted.org/packages/7b/07/a212ca4cfe56fa35c8315307e37df218e5946c726a792b5c7a795b245c10/dllist-2.0.0-py3-none-any.whl
+  name: dllist
+  version: 2.0.0
+  sha256: cd307b1a91bc46fae084f8c817d79be7e34951b149a2fd69004772e03573bfb3
+  requires_dist:
+  - pytest ; extra == 'test'
+  - pytest-cov ; extra == 'test'
+  requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/epoxy-1.5.10-h166bdaf_1.tar.bz2
+  sha256: 1e58ee2ed0f4699be202f23d49b9644b499836230da7dd5b2f63e6766acff89e
+  md5: a089d06164afd2d511347d3f87214e0b
+  depends:
+  - libgcc-ng >=10.3.0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 1440699
+  timestamp: 1648505042260
+- pypi: https://files.pythonhosted.org/packages/20/2c/90dd8d90acae15cab7ad344398030a0ba6a1cb1257d5e7da57f41e566940/equinox-0.11.11-py3-none-any.whl
+  name: equinox
+  version: 0.11.11
+  sha256: 49e9674f9bff0cde7ebcfbf2cdf4585c9231eb377eda31168bbf6467f88241e5
+  requires_dist:
+  - jax>=0.4.38
+  - jaxtyping>=0.2.20
+  - typing-extensions>=4.5.0
+  requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/esbuild-0.24.2-ha770c72_0.conda
+  sha256: f3e5eef234c5eda3f7841790e72e2bfdf8c19cd5fc2ff452c1015445d8eec921
+  md5: ad079d8d346d153ac11600626fe1ea38
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 3640554
+  timestamp: 1736588042523
+- pypi: https://files.pythonhosted.org/packages/05/83/bb4a4518bfa32a160dc455d8d944a6e00a7eb6759f449cb20c7b7879090e/etils-1.12.0-py3-none-any.whl
+  name: etils
+  version: 1.12.0
+  sha256: f80c2ff4289cc504b58b7e7a9f9db8b373a33227e43694a66808bcc81e51ffb8
+  requires_dist:
+  - etils[array-types] ; extra == 'all'
+  - etils[eapp] ; extra == 'all'
+  - etils[ecolab] ; extra == 'all'
+  - etils[edc] ; extra == 'all'
+  - etils[enp] ; extra == 'all'
+  - etils[epath] ; extra == 'all'
+  - etils[epath-gcs] ; extra == 'all'
+  - etils[epath-s3] ; extra == 'all'
+  - etils[epy] ; extra == 'all'
+  - etils[etqdm] ; extra == 'all'
+  - etils[etree] ; extra == 'all'
+  - etils[etree-dm] ; extra == 'all'
+  - etils[etree-jax] ; extra == 'all'
+  - etils[etree-tf] ; extra == 'all'
+  - etils[enp] ; extra == 'array-types'
+  - pytest ; extra == 'dev'
+  - pytest-subtests ; extra == 'dev'
+  - pytest-xdist ; extra == 'dev'
+  - pyink ; extra == 'dev'
+  - pylint>=2.6.0 ; extra == 'dev'
+  - chex ; extra == 'dev'
+  - fiddle ; extra == 'dev'
+  - torch ; extra == 'dev'
+  - optree ; extra == 'dev'
+  - tensorflow-datasets ; extra == 'dev'
+  - pydantic ; extra == 'dev'
+  - sphinx-apitree[ext] ; extra == 'docs'
+  - etils[dev,all] ; extra == 'docs'
+  - absl-py ; extra == 'eapp'
+  - simple-parsing ; extra == 'eapp'
+  - etils[epy] ; extra == 'eapp'
+  - jupyter ; extra == 'ecolab'
+  - numpy ; extra == 'ecolab'
+  - mediapy ; extra == 'ecolab'
+  - packaging ; extra == 'ecolab'
+  - protobuf ; extra == 'ecolab'
+  - etils[enp] ; extra == 'ecolab'
+  - etils[epy] ; extra == 'ecolab'
+  - etils[etree] ; extra == 'ecolab'
+  - etils[epy] ; extra == 'edc'
+  - numpy ; extra == 'enp'
+  - etils[epy] ; extra == 'enp'
+  - fsspec ; extra == 'epath'
+  - importlib-resources ; extra == 'epath'
+  - typing-extensions ; extra == 'epath'
+  - zipp ; extra == 'epath'
+  - etils[epy] ; extra == 'epath'
+  - gcsfs ; extra == 'epath-gcs'
+  - etils[epath] ; extra == 'epath-gcs'
+  - s3fs ; extra == 'epath-s3'
+  - etils[epath] ; extra == 'epath-s3'
+  - typing-extensions ; extra == 'epy'
+  - absl-py ; extra == 'etqdm'
+  - tqdm ; extra == 'etqdm'
+  - etils[epy] ; extra == 'etqdm'
+  - etils[array-types] ; extra == 'etree'
+  - etils[epy] ; extra == 'etree'
+  - etils[enp] ; extra == 'etree'
+  - etils[etqdm] ; extra == 'etree'
+  - dm-tree ; extra == 'etree-dm'
+  - etils[etree] ; extra == 'etree-dm'
+  - jax[cpu] ; extra == 'etree-jax'
+  - etils[etree] ; extra == 'etree-jax'
+  - tensorflow ; extra == 'etree-tf'
+  - etils[etree] ; extra == 'etree-tf'
+  - etils[ecolab] ; extra == 'lazy-imports'
+  requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/etuples-0.3.9-pyhd8ed1ab_1.conda
+  sha256: a2eb1d51f46b372bf1f514975b78c5492b431749bc86709969732c313bc2988e
+  md5: f2a2e0b6f6b043bcfa812408aa48a241
+  depends:
+  - cons
+  - multipledispatch
+  - python >=3.9
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/etuples?source=hash-mapping
+  size: 17573
+  timestamp: 1734526891894
+- conda: https://conda.anaconda.org/conda-forge/noarch/exceptiongroup-1.2.2-pyhd8ed1ab_1.conda
+  sha256: cbde2c64ec317118fc06b223c5fd87c8a680255e7348dd60e7b292d2e103e701
+  md5: a16662747cdeb9abbac74d0057cc976e
+  depends:
+  - python >=3.9
+  license: MIT and PSF-2.0
+  purls:
+  - pkg:pypi/exceptiongroup?source=hash-mapping
+  size: 20486
+  timestamp: 1733208916977
+- conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.1.0-pyhd8ed1ab_1.conda
+  sha256: 28d25ea375ebab4bf7479228f8430db20986187b04999136ff5c722ebd32eb60
+  md5: ef8b5fca76806159fc25b4f48d8737eb
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/executing?source=hash-mapping
+  size: 28348
+  timestamp: 1733569440265
+- conda: https://conda.anaconda.org/conda-forge/linux-64/expat-2.6.4-h5888daf_0.conda
+  sha256: 1848c7db9e264e3b8036ee133d570dd880422983cd20dd9585a505289606d276
+  md5: 1d6afef758879ef5ee78127eb4cd2c4a
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libexpat 2.6.4 h5888daf_0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 138145
+  timestamp: 1730967050578
+- conda: https://conda.anaconda.org/conda-forge/noarch/fasteners-0.19-pyhd8ed1ab_1.conda
+  sha256: 42fb170778b47303e82eddfea9a6d1e1b8af00c927cd5a34595eaa882b903a16
+  md5: dbe9d42e94b5ff7af7b7893f4ce052e7
+  depends:
+  - python >=3.9
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/fasteners?source=hash-mapping
+  size: 20711
+  timestamp: 1734943237791
+- conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.17.0-pyhd8ed1ab_0.conda
+  sha256: 006d7e5a0c17a6973596dd86bfc80d74ce541144d2aee2d22d46fd41df560a63
+  md5: 7f402b4a1007ee355bc50ce4d24d4a57
+  depends:
+  - python >=3.9
+  license: Unlicense
+  purls:
+  - pkg:pypi/filelock?source=hash-mapping
+  size: 17544
+  timestamp: 1737517924333
+- pypi: git+https://github.com/aseyboldt/flowjax?rev=07e7e32217bcfcaa7d68a304f332c925d26ab76f#07e7e32217bcfcaa7d68a304f332c925d26ab76f
+  name: flowjax
+  version: 17.0.2
+  requires_dist:
+  - jax>=0.4.16
+  - equinox>=0.10
+  - jaxtyping
+  - tqdm
+  - optax
+  - paramax
+  - pytest ; extra == 'dev'
+  - beartype ; extra == 'dev'
+  - ruff ; extra == 'dev'
+  - sphinx ; extra == 'dev'
+  - sphinx-book-theme ; extra == 'dev'
+  - sphinx-copybutton ; extra == 'dev'
+  - sphinx-autodoc-typehints ; extra == 'dev'
+  - nbsphinx ; extra == 'dev'
+  - ipython ; extra == 'dev'
+  - numpyro ; extra == 'dev'
+  requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-dejavu-sans-mono-2.37-hab24e00_0.tar.bz2
+  sha256: 58d7f40d2940dd0a8aa28651239adbf5613254df0f75789919c4e6762054403b
+  md5: 0c96522c6bdaed4b1566d11387caaf45
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 397370
+  timestamp: 1566932522327
+- conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-inconsolata-3.000-h77eed37_0.tar.bz2
+  sha256: c52a29fdac682c20d252facc50f01e7c2e7ceac52aa9817aaf0bb83f7559ec5c
+  md5: 34893075a5c9e55cdafac56607368fc6
+  license: OFL-1.1
+  license_family: Other
+  purls: []
+  size: 96530
+  timestamp: 1620479909603
+- conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-source-code-pro-2.038-h77eed37_0.tar.bz2
+  sha256: 00925c8c055a2275614b4d983e1df637245e19058d79fc7dd1a93b8d9fb4b139
+  md5: 4d59c254e01d9cde7957100457e2d5fb
+  license: OFL-1.1
+  license_family: Other
+  purls: []
+  size: 700814
+  timestamp: 1620479612257
+- conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-ubuntu-0.83-h77eed37_3.conda
+  sha256: 2821ec1dc454bd8b9a31d0ed22a7ce22422c0aef163c59f49dfdf915d0f0ca14
+  md5: 49023d73832ef61042f6a237cb2687e7
+  license: LicenseRef-Ubuntu-Font-Licence-Version-1.0
+  license_family: Other
+  purls: []
+  size: 1620504
+  timestamp: 1727511233259
+- conda: https://conda.anaconda.org/conda-forge/linux-64/fontconfig-2.15.0-h7e30c49_1.conda
+  sha256: 7093aa19d6df5ccb6ca50329ef8510c6acb6b0d8001191909397368b65b02113
+  md5: 8f5b0b297b59e1ac160ad4beec99dbee
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - freetype >=2.12.1,<3.0a0
+  - libexpat >=2.6.3,<3.0a0
+  - libgcc >=13
+  - libuuid >=2.38.1,<3.0a0
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 265599
+  timestamp: 1730283881107
+- conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-ecosystem-1-0.tar.bz2
+  sha256: a997f2f1921bb9c9d76e6fa2f6b408b7fa549edd349a77639c9fe7a23ea93e61
+  md5: fee5683a3f04bd15cbd8318b096a27ab
+  depends:
+  - fonts-conda-forge
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 3667
+  timestamp: 1566974674465
+- conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-forge-1-0.tar.bz2
+  sha256: 53f23a3319466053818540bcdf2091f253cbdbab1e0e9ae7b9e509dcaa2a5e38
+  md5: f766549260d6815b0c52253f1fb1bb29
+  depends:
+  - font-ttf-dejavu-sans-mono
+  - font-ttf-inconsolata
+  - font-ttf-source-code-pro
+  - font-ttf-ubuntu
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 4102
+  timestamp: 1566932280397
+- conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.56.0-py312h178313f_0.conda
+  sha256: 76ca95b4111fe27e64d74111b416b3462ad3db99f7109cbdf50e6e4b67dcf5b7
+  md5: 2f8a66f2f9eb931cdde040d02c6ab54c
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - brotli
+  - libgcc >=13
+  - munkres
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - unicodedata2 >=15.1.0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/fonttools?source=hash-mapping
+  size: 2834054
+  timestamp: 1738940929849
+- conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.12.1-h267a509_2.conda
+  sha256: b2e3c449ec9d907dd4656cb0dc93e140f447175b125a3824b31368b06c666bb6
+  md5: 9ae35c3d96db2c94ce0cef86efdfa2cb
+  depends:
+  - libgcc-ng >=12
+  - libpng >=1.6.39,<1.7.0a0
+  - libzlib >=1.2.13,<2.0.0a0
+  arch: x86_64
+  platform: linux
+  license: GPL-2.0-only OR FTL
+  purls: []
+  size: 634972
+  timestamp: 1694615932610
+- conda: https://conda.anaconda.org/conda-forge/linux-64/fribidi-1.0.10-h36c2ea0_0.tar.bz2
+  sha256: 5d7b6c0ee7743ba41399e9e05a58ccc1cfc903942e49ff6f677f6e423ea7a627
+  md5: ac7bc6a654f8f41b352b38f4051135f8
+  depends:
+  - libgcc-ng >=7.5.0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1
+  purls: []
+  size: 114383
+  timestamp: 1604416621168
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gcc-13.3.0-h9576a4e_1.conda
+  sha256: d0161362430183cbdbc3db9cf95f9a1af1793027f3ab8755b3d3586deb28bf84
+  md5: 606924335b5bcdf90e9aed9a2f5d22ed
+  depends:
+  - gcc_impl_linux-64 13.3.0.*
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 53864
+  timestamp: 1724801360210
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_impl_linux-64-13.3.0-hfea6d02_1.conda
+  sha256: 998ade1d487e93fc8a7a16b90e2af69ebb227355bf4646488661f7ae5887873c
+  md5: 0d043dbc126b64f79d915a0e96d3a1d5
+  depends:
+  - binutils_impl_linux-64 >=2.40
+  - libgcc >=13.3.0
+  - libgcc-devel_linux-64 13.3.0 h84ea5a7_101
+  - libgomp >=13.3.0
+  - libsanitizer 13.3.0 heb74ff8_1
+  - libstdcxx >=13.3.0
+  - sysroot_linux-64
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 67464415
+  timestamp: 1724801227937
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_linux-64-13.3.0-hc28eda2_7.conda
+  sha256: 1e5ac50580a68fdc7d2f5722abcf1a87898c24b1ab6eb5ecd322634742d93645
+  md5: ac23afbf5805389eb771e2ad3b476f75
+  depends:
+  - binutils_linux-64
+  - gcc_impl_linux-64 13.3.0.*
+  - sysroot_linux-64
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 32005
+  timestamp: 1731939593317
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gdk-pixbuf-2.42.12-hb9ae30d_0.conda
+  sha256: d5283b95a8d49dcd88d29b360d8b38694aaa905d968d156d72ab71d32b38facb
+  md5: 201db6c2d9a3c5e46573ac4cb2e92f4f
+  depends:
+  - libgcc-ng >=12
+  - libglib >=2.80.2,<3.0a0
+  - libjpeg-turbo >=3.0.0,<4.0a0
+  - libpng >=1.6.43,<1.7.0a0
+  - libtiff >=4.6.0,<4.8.0a0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  license_family: LGPL
+  purls: []
+  size: 528149
+  timestamp: 1715782983957
+- conda: https://conda.anaconda.org/conda-forge/linux-64/glib-tools-2.82.2-h4833e2c_1.conda
+  sha256: 5d8a48abdb1bc2b54f1380d2805cb9cd6cd9609ed0e5c3ed272aef92ab53b190
+  md5: e2e44caeaef6e4b107577aa46c95eb12
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libglib 2.82.2 h2ff4ddf_1
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 115452
+  timestamp: 1737037532892
+- conda: https://conda.anaconda.org/conda-forge/linux-64/graphite2-1.3.13-h59595ed_1003.conda
+  sha256: 0595b009f20f8f60f13a6398e7cdcbd2acea5f986633adcf85f5a2283c992add
+  md5: f87c7b7c2cb45f323ffbce941c78ab7c
+  depends:
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.0-or-later
+  license_family: LGPL
+  purls: []
+  size: 96855
+  timestamp: 1711634169756
+- conda: https://conda.anaconda.org/conda-forge/linux-64/graphviz-12.2.1-h5ae0cbf_1.conda
+  sha256: e6866409ba03df392ac5ec6f0d6ff9751a685ed917bfbcd8a73f550c5fe83c2b
+  md5: df7835d2c73cd1889d377cfd6694ada4
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - adwaita-icon-theme
+  - cairo >=1.18.2,<2.0a0
+  - fonts-conda-ecosystem
+  - gdk-pixbuf >=2.42.12,<3.0a0
+  - gtk3 >=3.24.43,<4.0a0
+  - gts >=0.7.6,<0.8.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libgd >=2.3.3,<2.4.0a0
+  - libglib >=2.82.2,<3.0a0
+  - librsvg >=2.58.4,<3.0a0
+  - libstdcxx >=13
+  - libwebp-base >=1.5.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - pango >=1.56.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: EPL-1.0
+  license_family: Other
+  purls: []
+  size: 2413095
+  timestamp: 1738602910851
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gtk3-3.24.43-h021d004_3.conda
+  sha256: c8f939497b43d90fa2ac9d99b44ed25759a798c305237300508e526de5e78de7
+  md5: 56c679bcdb8c1d824e927088725862cb
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - at-spi2-atk >=2.38.0,<3.0a0
+  - atk-1.0 >=2.38.0
+  - cairo >=1.18.2,<2.0a0
+  - epoxy >=1.5.10,<1.6.0a0
+  - fontconfig >=2.15.0,<3.0a0
+  - fonts-conda-ecosystem
+  - fribidi >=1.0.10,<2.0a0
+  - gdk-pixbuf >=2.42.12,<3.0a0
+  - glib-tools
+  - harfbuzz >=10.2.0,<11.0a0
+  - hicolor-icon-theme
+  - libcups >=2.3.3,<2.4.0a0
+  - libcups >=2.3.3,<3.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libglib >=2.82.2,<3.0a0
+  - liblzma >=5.6.3,<6.0a0
+  - libxkbcommon >=1.7.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - pango >=1.56.0,<2.0a0
+  - wayland >=1.23.1,<2.0a0
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxcomposite >=0.4.6,<1.0a0
+  - xorg-libxcursor >=1.2.3,<2.0a0
+  - xorg-libxdamage >=1.1.6,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxfixes >=6.0.1,<7.0a0
+  - xorg-libxi >=1.8.2,<2.0a0
+  - xorg-libxinerama >=1.1.5,<1.2.0a0
+  - xorg-libxrandr >=1.5.4,<2.0a0
+  - xorg-libxrender >=0.9.12,<0.10.0a0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.0-or-later
+  license_family: LGPL
+  purls: []
+  size: 5565328
+  timestamp: 1737497685605
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gts-0.7.6-h977cf35_4.conda
+  sha256: b5cd16262fefb836f69dc26d879b6508d29f8a5c5948a966c47fe99e2e19c99b
+  md5: 4d8df0b0db060d33c9a702ada998a8fe
+  depends:
+  - libgcc-ng >=12
+  - libglib >=2.76.3,<3.0a0
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.0-or-later
+  license_family: LGPL
+  purls: []
+  size: 318312
+  timestamp: 1686545244763
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gxx-13.3.0-h9576a4e_1.conda
+  sha256: 5446f5d1d609d996579f706d2020e83ef48e086d943bfeef7ab807ea246888a0
+  md5: 209182ca6b20aeff62f442e843961d81
+  depends:
+  - gcc 13.3.0.*
+  - gxx_impl_linux-64 13.3.0.*
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 53338
+  timestamp: 1724801498389
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gxx_impl_linux-64-13.3.0-hdbfa832_1.conda
+  sha256: 746dff24bb1efc89ab0ec108838d0711683054e3bbbcb94d042943410a98eca1
+  md5: 806367e23a0a6ad21e51875b34c57d7e
+  depends:
+  - gcc_impl_linux-64 13.3.0 hfea6d02_1
+  - libstdcxx-devel_linux-64 13.3.0 h84ea5a7_101
+  - sysroot_linux-64
+  - tzdata
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 13337720
+  timestamp: 1724801455825
+- conda: https://conda.anaconda.org/conda-forge/linux-64/gxx_linux-64-13.3.0-h6834431_7.conda
+  sha256: a9b1ffea76f2cc5aedeead4793fcded7a687cce9d5e3f4fe93629f1b1d5043a6
+  md5: 7c82ca9bda609b6f72f670e4219d3787
+  depends:
+  - binutils_linux-64
+  - gcc_linux-64 13.3.0 hc28eda2_7
+  - gxx_impl_linux-64 13.3.0.*
+  - sysroot_linux-64
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 30356
+  timestamp: 1731939612705
+- conda: https://conda.anaconda.org/conda-forge/noarch/h5netcdf-1.5.0-pyhd8ed1ab_0.conda
+  sha256: 2dd1aa54eb0b0f0e15db77d6dd3f16e532903f99a0375283142dd3df4d990e46
+  md5: af8ab1ff0815078c40ba96f47f48f353
+  depends:
+  - h5py
+  - packaging
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/h5netcdf?source=hash-mapping
+  size: 46375
+  timestamp: 1737887368948
+- conda: https://conda.anaconda.org/conda-forge/linux-64/h5py-3.13.0-nompi_py312hedeef09_100.conda
+  sha256: 76bb853325f0c756599edb0be014723b01fea61e24817fd2f0b9ddfe4c570c0f
+  md5: ed73cf6f5e1ce5e823e6efcf54cbdc51
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - cached-property
+  - hdf5 >=1.14.3,<1.14.4.0a0
+  - libgcc >=13
+  - numpy >=1.19,<3
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/h5py?source=hash-mapping
+  size: 1388165
+  timestamp: 1739952623855
+- conda: https://conda.anaconda.org/conda-forge/linux-64/harfbuzz-10.2.0-h4bba637_0.conda
+  sha256: 94426eca8c60b43f57beb3338d3298dda09452c7a42314bbbb4ebfa552542a84
+  md5: 9e38e86167e8b1ea0094747d12944ce4
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - cairo >=1.18.2,<2.0a0
+  - freetype >=2.12.1,<3.0a0
+  - graphite2
+  - icu >=75.1,<76.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libglib >=2.82.2,<3.0a0
+  - libstdcxx >=13
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 1646987
+  timestamp: 1736702906600
+- conda: https://conda.anaconda.org/conda-forge/linux-64/hdf5-1.14.3-nompi_h2d575fe_109.conda
+  sha256: e8669a6d76d415f4fdbe682507ac3a3b39e8f493d2f2bdc520817f80b7cc0753
+  md5: e7a7a6e6f70553a31e6e79c65768d089
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libaec >=1.1.3,<2.0a0
+  - libcurl >=8.11.1,<9.0a0
+  - libgcc >=13
+  - libgfortran
+  - libgfortran5 >=13.3.0
+  - libstdcxx >=13
+  - libzlib >=1.3.1,<2.0a0
+  - openssl >=3.4.0,<4.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 3930078
+  timestamp: 1737516601132
+- conda: https://conda.anaconda.org/conda-forge/linux-64/hicolor-icon-theme-0.17-ha770c72_2.tar.bz2
+  sha256: 336f29ceea9594f15cc8ec4c45fdc29e10796573c697ee0d57ebb7edd7e92043
+  md5: bbf6f174dcd3254e19a2f5d2295ce808
+  arch: x86_64
+  platform: linux
+  license: GPL-2.0-or-later
+  license_family: GPL
+  purls: []
+  size: 13841
+  timestamp: 1605162808667
+- conda: https://conda.anaconda.org/conda-forge/linux-64/icu-75.1-he02047a_0.conda
+  sha256: 71e750d509f5fa3421087ba88ef9a7b9be11c53174af3aa4d06aff4c18b38e8e
+  md5: 8b189310083baabfb622af68fd9d3ae3
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 12129203
+  timestamp: 1720853576813
+- pypi: https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl
+  name: idna
+  version: '3.10'
+  sha256: 946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3
+  requires_dist:
+  - ruff>=0.6.2 ; extra == 'all'
+  - mypy>=1.11.2 ; extra == 'all'
+  - pytest>=8.3.2 ; extra == 'all'
+  - flake8>=7.1.1 ; extra == 'all'
+  requires_python: '>=3.6'
+- conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
+  sha256: 598951ebdb23e25e4cec4bbff0ae369cec65ead80b50bc08b441d8e54de5cf03
+  md5: f4b39bf00c69f56ac01e020ebfac066c
+  depends:
+  - python >=3.9
+  - zipp >=0.5
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/importlib-metadata?source=hash-mapping
+  size: 29141
+  timestamp: 1737420302391
+- conda: https://conda.anaconda.org/conda-forge/noarch/importlib_resources-6.5.2-pyhd8ed1ab_0.conda
+  sha256: acc1d991837c0afb67c75b77fdc72b4bf022aac71fedd8b9ea45918ac9b08a80
+  md5: c85c76dc67d75619a92f51dfbce06992
+  depends:
+  - python >=3.9
+  - zipp >=3.1.0
+  constrains:
+  - importlib-resources >=6.5.2,<6.5.3.0a0
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/importlib-resources?source=hash-mapping
+  size: 33781
+  timestamp: 1736252433366
+- conda: https://conda.anaconda.org/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
+  sha256: 0ec8f4d02053cd03b0f3e63168316530949484f80e16f5e2fb199a1d117a89ca
+  md5: 6837f3eff7dcea42ecd714ce1ac2b108
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/iniconfig?source=hash-mapping
+  size: 11474
+  timestamp: 1733223232820
+- conda: https://conda.anaconda.org/conda-forge/noarch/ipykernel-6.29.5-pyh3099207_0.conda
+  sha256: 33cfd339bb4efac56edf93474b37ddc049e08b1b4930cf036c893cc1f5a1f32a
+  md5: b40131ab6a36ac2c09b7c57d4d3fbf99
+  depends:
+  - __linux
+  - comm >=0.1.1
+  - debugpy >=1.6.5
+  - ipython >=7.23.1
+  - jupyter_client >=6.1.12
+  - jupyter_core >=4.12,!=5.0.*
+  - matplotlib-inline >=0.1
+  - nest-asyncio
+  - packaging
+  - psutil
+  - python >=3.8
+  - pyzmq >=24
+  - tornado >=6.1
+  - traitlets >=5.4.0
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/ipykernel?source=hash-mapping
+  size: 119084
+  timestamp: 1719845605084
+- conda: https://conda.anaconda.org/conda-forge/noarch/ipython-8.32.0-pyh907856f_0.conda
+  sha256: b1b940cfe85d5f0aaed83ef8c9f07ee80daa68acb05feeb5142d620472b01e0d
+  md5: 9de86472b8f207fb098c69daaad50e67
+  depends:
+  - __unix
+  - pexpect >4.3
+  - python >=3.10
+  - decorator
+  - exceptiongroup
+  - jedi >=0.16
+  - matplotlib-inline
+  - pickleshare
+  - prompt-toolkit >=3.0.41,<3.1.0
+  - pygments >=2.4.0
+  - stack_data
+  - traitlets >=5.13.0
+  - typing_extensions >=4.6
+  - python
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/ipython?source=hash-mapping
+  size: 636676
+  timestamp: 1738421264236
+- conda: https://conda.anaconda.org/conda-forge/noarch/ipywidgets-8.1.5-pyhd8ed1ab_1.conda
+  sha256: f419657566e3d9bea85b288a0ce3a8e42d76cd82ac1697c6917891df3ae149ab
+  md5: bb19ad65196475ab6d0bb3532d7f8d96
+  depends:
+  - comm >=0.1.3
+  - ipython >=6.1.0
+  - jupyterlab_widgets >=3.0.13,<3.1.0
+  - python >=3.9
+  - traitlets >=4.3.1
+  - widgetsnbextension >=4.0.13,<4.1.0
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/ipywidgets?source=hash-mapping
+  size: 113982
+  timestamp: 1733493669268
+- pypi: https://files.pythonhosted.org/packages/f4/58/cc0721a1030fcbab0984beea0bf3c4610ec103f738423cdfa9c4ceb40598/jax-0.5.0-py3-none-any.whl
+  name: jax
+  version: 0.5.0
+  sha256: b3907aa87ae2c340b39cdbf80c07a74550369cafcaf7398fb60ba58d167345ab
+  requires_dist:
+  - jaxlib<=0.5.0,>=0.5.0
+  - ml-dtypes>=0.4.0
+  - numpy>=1.25
+  - numpy>=1.26.0 ; python_full_version >= '3.12'
+  - opt-einsum
+  - scipy>=1.11.1
+  - jaxlib==0.5.0 ; extra == 'minimum-jaxlib'
+  - jaxlib==0.4.38 ; extra == 'ci'
+  - jaxlib<=0.5.0,>=0.5.0 ; extra == 'tpu'
+  - libtpu-nightly==0.1.dev20241010+nightly.cleanup ; extra == 'tpu'
+  - libtpu==0.0.8 ; extra == 'tpu'
+  - requests ; extra == 'tpu'
+  - jaxlib==0.5.0 ; extra == 'cuda'
+  - jax-cuda12-plugin[with-cuda]<=0.5.0,>=0.5.0 ; extra == 'cuda'
+  - jaxlib==0.5.0 ; extra == 'cuda12'
+  - jax-cuda12-plugin[with-cuda]<=0.5.0,>=0.5.0 ; extra == 'cuda12'
+  - jaxlib==0.5.0 ; extra == 'cuda12-pip'
+  - jax-cuda12-plugin[with-cuda]<=0.5.0,>=0.5.0 ; extra == 'cuda12-pip'
+  - jaxlib==0.5.0 ; extra == 'cuda12-local'
+  - jax-cuda12-plugin==0.5.0 ; extra == 'cuda12-local'
+  - jaxlib==0.5.0 ; extra == 'rocm'
+  - jax-rocm60-plugin<=0.5.0,>=0.5.0 ; extra == 'rocm'
+  - kubernetes ; extra == 'k8s'
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/66/e9/211ba3e46ec22c722c4d61a739cfccf79b0618006d6f5fa53eb4eb93ed6d/jaxlib-0.5.0-cp312-cp312-manylinux2014_x86_64.whl
+  name: jaxlib
+  version: 0.5.0
+  sha256: f980c733e98c998a8da87c9a8cc61b6726d0be667a58bd664c1d717b4b4eae75
+  requires_dist:
+  - scipy>=1.11.1
+  - numpy>=1.25
+  - ml-dtypes>=0.2.0
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/db/7e/da7b57a1f3af7303a0f3c8594d820fc0d3a9bbe3810a357eb21eb166e76b/jaxtyping-0.2.38-py3-none-any.whl
+  name: jaxtyping
+  version: 0.2.38
+  sha256: bc209ab8ec29917b6f0c7dec4a8ea1fc276f7d94f25b71c01d1243ec2b21ae12
+  requires_dist:
+  - wadler-lindig>=0.1.3
+  requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
+  sha256: 92c4d217e2dc68983f724aa983cca5464dcb929c566627b26a2511159667dba8
+  md5: a4f4c5dc9b80bc50e0d3dc4e6e8f1bd9
+  depends:
+  - parso >=0.8.3,<0.9.0
+  - python >=3.9
+  license: Apache-2.0 AND MIT
+  purls:
+  - pkg:pypi/jedi?source=hash-mapping
+  size: 843646
+  timestamp: 1733300981994
+- conda: https://conda.anaconda.org/conda-forge/noarch/jsonschema-4.23.0-pyhd8ed1ab_1.conda
+  sha256: be992a99e589146f229c58fe5083e0b60551d774511c494f91fe011931bd7893
+  md5: a3cead9264b331b32fe8f0aabc967522
+  depends:
+  - attrs >=22.2.0
+  - importlib_resources >=1.4.0
+  - jsonschema-specifications >=2023.03.6
+  - pkgutil-resolve-name >=1.3.10
+  - python >=3.9
+  - referencing >=0.28.4
+  - rpds-py >=0.7.1
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/jsonschema?source=hash-mapping
+  size: 74256
+  timestamp: 1733472818764
+- conda: https://conda.anaconda.org/conda-forge/noarch/jsonschema-specifications-2024.10.1-pyhd8ed1ab_1.conda
+  sha256: 37127133837444cf0e6d1a95ff5a505f8214ed4e89e8e9343284840e674c6891
+  md5: 3b519bc21bc80e60b456f1e62962a766
+  depends:
+  - python >=3.9
+  - referencing >=0.31.0
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/jsonschema-specifications?source=hash-mapping
+  size: 16170
+  timestamp: 1733493624968
+- conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_client-8.6.3-pyhd8ed1ab_1.conda
+  sha256: 19d8bd5bb2fde910ec59e081eeb59529491995ce0d653a5209366611023a0b3a
+  md5: 4ebae00eae9705b0c3d6d1018a81d047
+  depends:
+  - importlib-metadata >=4.8.3
+  - jupyter_core >=4.12,!=5.0.*
+  - python >=3.9
+  - python-dateutil >=2.8.2
+  - pyzmq >=23.0
+  - tornado >=6.2
+  - traitlets >=5.3
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/jupyter-client?source=hash-mapping
+  size: 106342
+  timestamp: 1733441040958
+- conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_core-5.7.2-pyh31011fe_1.conda
+  sha256: 732b1e8536bc22a5a174baa79842d79db2f4956d90293dd82dc1b3f6099bcccd
+  md5: 0a2980dada0dd7fd0998f0342308b1b1
+  depends:
+  - __unix
+  - platformdirs >=2.5
+  - python >=3.8
+  - traitlets >=5.3
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/jupyter-core?source=hash-mapping
+  size: 57671
+  timestamp: 1727163547058
+- conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_widgets-3.0.13-pyhd8ed1ab_1.conda
+  sha256: 206489e417408d2ffc2a7b245008b4735a8beb59df6c9109d4f77e7bc5969d5d
+  md5: b26e487434032d7f486277beb0cead3a
+  depends:
+  - python >=3.9
+  constrains:
+  - jupyterlab >=3,<5
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/jupyterlab-widgets?source=hash-mapping
+  size: 186358
+  timestamp: 1733428156991
+- conda: https://conda.anaconda.org/conda-forge/noarch/kernel-headers_linux-64-3.10.0-he073ed8_18.conda
+  sha256: a922841ad80bd7b222502e65c07ecb67e4176c4fa5b03678a005f39fcc98be4b
+  md5: ad8527bf134a90e1c9ed35fa0b64318c
+  constrains:
+  - sysroot_linux-64 ==2.17
+  license: LGPL-2.0-or-later AND LGPL-2.0-or-later WITH exceptions AND GPL-2.0-or-later AND MPL-2.0
+  license_family: GPL
+  purls: []
+  size: 943486
+  timestamp: 1729794504440
+- conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.1-h166bdaf_0.tar.bz2
+  sha256: 150c05a6e538610ca7c43beb3a40d65c90537497a4f6a5f4d15ec0451b6f5ebb
+  md5: 30186d27e2c9fa62b45fb1476b7200e3
+  depends:
+  - libgcc-ng >=10.3.0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 117831
+  timestamp: 1646151697040
+- conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.8-py312h84d6215_0.conda
+  sha256: 3ce99d721c1543f6f8f5155e53eef11be47b2f5942a8d1060de6854f9d51f246
+  md5: 6713467dc95509683bfa3aca08524e8a
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/kiwisolver?source=hash-mapping
+  size: 71649
+  timestamp: 1736908364705
+- conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
+  sha256: 99df692f7a8a5c27cd14b5fb1374ee55e756631b9c3d659ed3ee60830249b238
+  md5: 3f43953b7d3fb3aaa1d0d0723d91e368
+  depends:
+  - keyutils >=1.6.1,<2.0a0
+  - libedit >=3.1.20191231,<3.2.0a0
+  - libedit >=3.1.20191231,<4.0a0
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  - openssl >=3.3.1,<4.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 1370023
+  timestamp: 1719463201255
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
+  sha256: d6a61830a354da022eae93fa896d0991385a875c6bba53c82263a289deda9db8
+  md5: 000e85703f0fd9594c81710dd5066471
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libjpeg-turbo >=3.0.0,<4.0a0
+  - libtiff >=4.7.0,<4.8.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 248046
+  timestamp: 1739160907615
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.43-h712a8e2_2.conda
+  sha256: 7c91cea91b13f4314d125d1bedb9d03a29ebbd5080ccdea70260363424646dbe
+  md5: 048b02e3962f066da18efe3a21b77672
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  constrains:
+  - binutils_impl_linux-64 2.43
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only
+  license_family: GPL
+  purls: []
+  size: 669211
+  timestamp: 1729655358674
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h27087fc_0.tar.bz2
+  sha256: cb55f36dcd898203927133280ae1dc643368af041a48bcf7c026acb7c47b0c12
+  md5: 76bbff344f0134279f225174e9064c8f
+  depends:
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls: []
+  size: 281798
+  timestamp: 1657977462600
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libaec-1.1.3-h59595ed_0.conda
+  sha256: 2ef420a655528bca9d269086cf33b7e90d2f54ad941b437fb1ed5eca87cee017
+  md5: 5e97e271911b8b2001a8b71860c32faa
+  depends:
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls: []
+  size: 35446
+  timestamp: 1711021212685
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-30_h66dfbfd_blis.conda
+  build_number: 30
+  sha256: 4b2743e2172a85eb984f3c167d6027295468b6c8a87b4445eb6551c55a1dbd23
+  md5: e566a0d2bf6341d06ce06b74f2f9dc17
+  depends:
+  - blis >=0.9.0,<0.9.1.0a0
+  - mkl >=2024.2.2,<2025.0a0
+  constrains:
+  - blas =2.130=blis
+  - libcblas =3.9.0=30*_blis
+  arch: x86_64
+  platform: linux
+  track_features:
+  - blas_blis
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 16949
+  timestamp: 1739836207638
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hb9d3cd8_2.conda
+  sha256: d9db2de60ea917298e658143354a530e9ca5f9c63471c65cf47ab39fd2f429e3
+  md5: 41b599ed2b02abcfdd84302bff174b23
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 68851
+  timestamp: 1725267660471
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hb9d3cd8_2.conda
+  sha256: 2892d512cad096cb03f1b66361deeab58b64e15ba525d6592bb6d609e7045edf
+  md5: 9566f0bd264fbd463002e759b8a82401
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libbrotlicommon 1.1.0 hb9d3cd8_2
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 32696
+  timestamp: 1725267669305
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hb9d3cd8_2.conda
+  sha256: 779f58174e99de3600e939fa46eddb453ec5d3c60bb46cdaa8b4c127224dbf29
+  md5: 06f70867945ea6a84d35836af780f1de
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libbrotlicommon 1.1.0 hb9d3cd8_2
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 281750
+  timestamp: 1725267679782
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-30_hba4ea11_blis.conda
+  build_number: 30
+  sha256: 48c142e0bec47c0d1f3c2d67373b91812058e78746565b70e388ec5ddde2366e
+  md5: addf030df95030b6c91f346d39995ce8
+  depends:
+  - libblas 3.9.0 30_h66dfbfd_blis
+  - mkl >=2024.2.2,<2025.0a0
+  constrains:
+  - blas =2.130=blis
+  arch: x86_64
+  platform: linux
+  track_features:
+  - blas_blis
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 16921
+  timestamp: 1739836217216
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcups-2.3.3-h4637d8d_4.conda
+  sha256: bc67b9b21078c99c6bd8595fe7e1ed6da1f721007726e717f0449de7032798c4
+  md5: d4529f4dff3057982a7617c7ac58fde3
+  depends:
+  - krb5 >=1.21.1,<1.22.0a0
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  - libzlib >=1.2.13,<2.0.0a0
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls: []
+  size: 4519402
+  timestamp: 1689195353551
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.12.1-h332b0f4_0.conda
+  sha256: 2ebc3039af29269e4cdb858fca36265e5e400c1125a4bcd84ae73a596e0e76ca
+  md5: 45e9dc4e7b25e2841deb392be085500e
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - krb5 >=1.21.3,<1.22.0a0
+  - libgcc >=13
+  - libnghttp2 >=1.64.0,<2.0a0
+  - libssh2 >=1.11.1,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - openssl >=3.4.1,<4.0a0
+  - zstd >=1.5.6,<1.6.0a0
+  arch: x86_64
+  platform: linux
+  license: curl
+  license_family: MIT
+  purls: []
+  size: 426675
+  timestamp: 1739512336799
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.23-h4ddbbb0_0.conda
+  sha256: 511d801626d02f4247a04fff957cc6e9ec4cc7e8622bd9acd076bcdc5de5fe66
+  md5: 8dfae1d2e74767e9ce36d5fa0d8605db
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 72255
+  timestamp: 1734373823254
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h7949ede_0.conda
+  sha256: d789471216e7aba3c184cd054ed61ce3f6dac6f87a50ec69291b9297f8c18724
+  md5: c277e0a4d549b03ac1e9d6cbbe3d017b
+  depends:
+  - ncurses
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - ncurses >=6.5,<7.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls: []
+  size: 134676
+  timestamp: 1738479519902
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-hd590300_2.conda
+  sha256: 1cd6048169fa0395af74ed5d8f1716e22c19a81a8a36f934c110ca3ad4dd27b4
+  md5: 172bf1cd1ff8629f2b1179945ed45055
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls: []
+  size: 112766
+  timestamp: 1702146165126
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.6.4-h5888daf_0.conda
+  sha256: 56541b98447b58e52d824bd59d6382d609e11de1f8adf20b23143e353d2b8d26
+  md5: db833e03127376d461e1e13e76f09b6c
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  constrains:
+  - expat 2.6.4.*
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 73304
+  timestamp: 1730967041968
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.4.6-h2dba641_0.conda
+  sha256: 67a6c95e33ebc763c1adc3455b9a9ecde901850eb2fceb8e646cc05ef3a663da
+  md5: e3eb7806380bc8bcecba6d749ad5f026
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 53415
+  timestamp: 1739260413716
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-14.2.0-h77fa898_1.conda
+  sha256: 53eb8a79365e58849e7b1a068d31f4f9e718dc938d6f2c03e960345739a03569
+  md5: 3cb76c3f10d3bc7f1105b2fc9db984df
+  depends:
+  - _libgcc_mutex 0.1 conda_forge
+  - _openmp_mutex >=4.5
+  constrains:
+  - libgomp 14.2.0 h77fa898_1
+  - libgcc-ng ==14.2.0=*_1
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 848745
+  timestamp: 1729027721139
+- conda: https://conda.anaconda.org/conda-forge/noarch/libgcc-devel_linux-64-13.3.0-h84ea5a7_101.conda
+  sha256: 027cfb011328a108bc44f512a2dec6d954db85709e0b79b748c3392f85de0c64
+  md5: 0ce69d40c142915ac9734bc6134e514a
+  depends:
+  - __unix
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 2598313
+  timestamp: 1724801050802
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-14.2.0-h69a702a_1.conda
+  sha256: 3a76969c80e9af8b6e7a55090088bc41da4cffcde9e2c71b17f44d37b7cb87f7
+  md5: e39480b9ca41323497b05492a63bc35b
+  depends:
+  - libgcc 14.2.0 h77fa898_1
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 54142
+  timestamp: 1729027726517
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgd-2.3.3-h6f5c62b_11.conda
+  sha256: 19e5be91445db119152217e8e8eec4fd0499d854acc7d8062044fb55a70971cd
+  md5: 68fc66282364981589ef36868b1a7c78
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - fontconfig >=2.15.0,<3.0a0
+  - fonts-conda-ecosystem
+  - freetype >=2.12.1,<3.0a0
+  - icu >=75.1,<76.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libjpeg-turbo >=3.0.0,<4.0a0
+  - libpng >=1.6.45,<1.7.0a0
+  - libtiff >=4.7.0,<4.8.0a0
+  - libwebp-base >=1.5.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: GD
+  license_family: BSD
+  purls: []
+  size: 177082
+  timestamp: 1737548051015
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-14.2.0-h69a702a_1.conda
+  sha256: fc9e7f22a17faf74da904ebfc4d88699013d2992e55505e4aa0eb01770290977
+  md5: f1fd30127802683586f768875127a987
+  depends:
+  - libgfortran5 14.2.0 hd5240d6_1
+  constrains:
+  - libgfortran-ng ==14.2.0=*_1
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 53997
+  timestamp: 1729027752995
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-14.2.0-hd5240d6_1.conda
+  sha256: d149a37ca73611e425041f33b9d8dbed6e52ec506fe8cc1fc0ee054bddeb6d5d
+  md5: 9822b874ea29af082e5d36098d25427d
+  depends:
+  - libgcc >=14.2.0
+  constrains:
+  - libgfortran 14.2.0
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 1462645
+  timestamp: 1729027735353
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libglib-2.82.2-h2ff4ddf_1.conda
+  sha256: f0804a9e46ae7b32ca698d26c1c95aa82a91f71b6051883d4a46bea725be9ea4
+  md5: 37d1af619d999ee8f1f73cf5a06f4e2f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libffi >=3.4,<4.0a0
+  - libgcc >=13
+  - libiconv >=1.17,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - pcre2 >=10.44,<10.45.0a0
+  constrains:
+  - glib 2.82.2 *_1
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 3923974
+  timestamp: 1737037491054
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-14.2.0-h77fa898_1.conda
+  sha256: 1911c29975ec99b6b906904040c855772ccb265a1c79d5d75c8ceec4ed89cd63
+  md5: cc3573974587f12dda90d96e3e55a702
+  depends:
+  - _libgcc_mutex 0.1 conda_forge
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 460992
+  timestamp: 1729027639220
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libhwloc-2.11.2-default_h0d58e46_1001.conda
+  sha256: d14c016482e1409ae1c50109a9ff933460a50940d2682e745ab1c172b5282a69
+  md5: 804ca9e91bcaea0824a341d55b1684f2
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - libxml2 >=2.13.4,<3.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 2423200
+  timestamp: 1731374922090
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h4ce23a2_0.conda
+  sha256: 659a83c0c184a30336ca512d670a7db56bfb5e289caa9ee641e6f78fd2326463
+  md5: 0743bcd51fc4e6e0367e0aaf41e9bebc
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-only
+  purls: []
+  size: 713745
+  timestamp: 1739866934640
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.0.0-hd590300_1.conda
+  sha256: b954e09b7e49c2f2433d6f3bb73868eda5e378278b0f8c1dd10a7ef090e14f2f
+  md5: ea25936bb4080d843790b586850f82b8
+  depends:
+  - libgcc-ng >=12
+  constrains:
+  - jpeg <0.0.0a
+  arch: x86_64
+  platform: linux
+  license: IJG AND BSD-3-Clause AND Zlib
+  purls: []
+  size: 618575
+  timestamp: 1694474974816
+- conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-9_h9f1adc1_netlib.conda
+  build_number: 9
+  sha256: b2d7d3ae78113afd8c147b91cc54bfdf7a3b1bcec7e9b2b562f3f9a3ba0edf7c
+  md5: e35381a8ed8c80a987dc804384c34aed
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libblas 3.9.0.*
+  - libgcc >=13
+  - libgfortran
+  - libgfortran5 >=13.3.0
+  - mkl >=2024.2.2,<2025.0a0
+  arch: x86_64
+  platform: linux
+  track_features:
+  - blas_netlib
+  - blas_netlib_2
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 2781524
+  timestamp: 1739843060105
+- conda: https://conda.anaconda.org/conda-forge/linux-64/liblapacke-3.9.0-9_h0ad7b2f_netlib.conda
+  build_number: 9
+  sha256: e0ac9d2e61819d77564f168706325b443c87480f4f240321aa6ef2bc60dd7acd
+  md5: df14f215234841c75debd3618a9e1e8c
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libblas 3.9.0.*
+  - libcblas 3.9.0.*
+  - libgcc >=13
+  - libgfortran
+  - libgfortran5 >=13.3.0
+  - liblapack 3.9.0.*
+  - mkl >=2024.2.2,<2025.0a0
+  arch: x86_64
+  platform: linux
+  track_features:
+  - blas_netlib
+  - blas_netlib_2
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 496051
+  timestamp: 1739843076862
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libllvm14-14.0.6-hcd5def8_4.conda
+  sha256: 225cc7c3b20ac1db1bdb37fa18c95bf8aecef4388e984ab2f7540a9f4382106a
+  md5: 73301c133ded2bf71906aa2104edae8b
+  depends:
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  - libzlib >=1.2.13,<2.0.0a0
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0 WITH LLVM-exception
+  license_family: Apache
+  purls: []
+  size: 31484415
+  timestamp: 1690557554081
+- conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.6.4-hb9d3cd8_0.conda
+  sha256: cad52e10319ca4585bc37f0bc7cce99ec7c15dc9168e42ccb96b741b0a27db3f
+  md5: 42d5b6a0f30d3c10cd88cb8584fda1cb
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: 0BSD
+  purls: []
+  size: 111357
+  timestamp: 1738525339684
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.64.0-h161d5f1_0.conda
+  sha256: b0f2b3695b13a989f75d8fd7f4778e1c7aabe3b36db83f0fe80b2cd812c0e975
+  md5: 19e57602824042dfd0446292ef90488b
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - c-ares >=1.32.3,<2.0a0
+  - libev >=4.33,<4.34.0a0
+  - libev >=4.33,<5.0a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - libzlib >=1.3.1,<2.0a0
+  - openssl >=3.3.2,<4.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 647599
+  timestamp: 1729571887612
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hd590300_0.conda
+  sha256: 26d77a3bb4dceeedc2a41bd688564fe71bf2d149fdcf117049970bc02ff1add6
+  md5: 30fd6e37fe21f86f4bd26d6ee73eeec7
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-only
+  license_family: GPL
+  purls: []
+  size: 33408
+  timestamp: 1697359010159
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.47-h943b412_0.conda
+  sha256: 23367d71da58c9a61c8cbd963fcffb92768d4ae5ffbef9a47cdf1f54f98c5c36
+  md5: 55199e2ae2c3651f6f9b2a447b47bdc9
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: zlib-acknowledgement
+  purls: []
+  size: 288701
+  timestamp: 1739952993639
+- conda: https://conda.anaconda.org/conda-forge/linux-64/librsvg-2.58.4-h49af25d_2.conda
+  sha256: 475013475a3209c24a82f9e80c545d56ccca2fa04df85952852f3d73caa38ff9
+  md5: b9846db0abffb09847e2cb0fec4b4db6
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - cairo >=1.18.2,<2.0a0
+  - freetype >=2.12.1,<3.0a0
+  - gdk-pixbuf >=2.42.12,<3.0a0
+  - harfbuzz >=10.1.0,<11.0a0
+  - libgcc >=13
+  - libglib >=2.82.2,<3.0a0
+  - libpng >=1.6.44,<1.7.0a0
+  - libxml2 >=2.13.5,<3.0a0
+  - pango >=1.54.0,<2.0a0
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 6342757
+  timestamp: 1734902068235
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libsanitizer-13.3.0-heb74ff8_1.conda
+  sha256: c86d130f0a3099e46ff51aa7ffaab73cb44fc420d27a96076aab3b9a326fc137
+  md5: c4cb22f270f501f5c59a122dc2adf20a
+  depends:
+  - libgcc >=13.3.0
+  - libstdcxx >=13.3.0
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 4133922
+  timestamp: 1724801171589
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libsodium-1.0.20-h4ab18f5_0.conda
+  sha256: 0105bd108f19ea8e6a78d2d994a6d4a8db16d19a41212070d2d1d48a63c34161
+  md5: a587892d3c13b6621a6091be690dbca2
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: ISC
+  purls: []
+  size: 205978
+  timestamp: 1716828628198
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.49.1-hee588c1_1.conda
+  sha256: 7a09eef804ef7cf4d88215c2297eabb72af8ad0bd5b012060111c289f14bbe7d
+  md5: 73cea06049cc4174578b432320a003b8
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: Unlicense
+  purls: []
+  size: 915956
+  timestamp: 1739953155793
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-hf672d98_0.conda
+  sha256: 0407ac9fda2bb67e11e357066eff144c845801d00b5f664efbc48813af1e7bb9
+  md5: be2de152d8073ef1c01b7728475f2fe7
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libzlib >=1.3.1,<2.0a0
+  - openssl >=3.4.0,<4.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 304278
+  timestamp: 1732349402869
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-14.2.0-hc0a3c3a_1.conda
+  sha256: 4661af0eb9bdcbb5fb33e5d0023b001ad4be828fccdcc56500059d56f9869462
+  md5: 234a5554c53625688d51062645337328
+  depends:
+  - libgcc 14.2.0 h77fa898_1
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 3893695
+  timestamp: 1729027746910
+- conda: https://conda.anaconda.org/conda-forge/noarch/libstdcxx-devel_linux-64-13.3.0-h84ea5a7_101.conda
+  sha256: 0a9226c1b994f996229ffb54fa40d608cd4e4b48e8dc73a66134bea8ce949412
+  md5: 29b5a4ed4613fa81a07c21045e3f5bf6
+  depends:
+  - __unix
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 14074676
+  timestamp: 1724801075448
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-14.2.0-h4852527_1.conda
+  sha256: 25bb30b827d4f6d6f0522cc0579e431695503822f144043b93c50237017fffd8
+  md5: 8371ac6457591af2cf6159439c1fd051
+  depends:
+  - libstdcxx 14.2.0 hc0a3c3a_1
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only WITH GCC-exception-3.1
+  license_family: GPL
+  purls: []
+  size: 54105
+  timestamp: 1729027780628
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.0-hd9ff511_3.conda
+  sha256: b224e16b88d76ea95e4af56e2bc638c603bd26a770b98d117d04541d3aafa002
+  md5: 0ea6510969e1296cc19966fad481f6de
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - lerc >=4.0.0,<5.0a0
+  - libdeflate >=1.23,<1.24.0a0
+  - libgcc >=13
+  - libjpeg-turbo >=3.0.0,<4.0a0
+  - liblzma >=5.6.3,<6.0a0
+  - libstdcxx >=13
+  - libwebp-base >=1.4.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - zstd >=1.5.6,<1.6.0a0
+  arch: x86_64
+  platform: linux
+  license: HPND
+  purls: []
+  size: 428173
+  timestamp: 1734398813264
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.38.1-h0b41bf4_0.conda
+  sha256: 787eb542f055a2b3de553614b25f09eefb0a0931b0c87dbcce6efdfd92f04f18
+  md5: 40b61aab5c7ba9ff276c41cfffe6b80b
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 33601
+  timestamp: 1680112270483
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.5.0-h851e524_0.conda
+  sha256: c45283fd3e90df5f0bd3dbcd31f59cdd2b001d424cf30a07223655413b158eaf
+  md5: 63f790534398730f59e1b899c3644d4a
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  constrains:
+  - libwebp 1.5.0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 429973
+  timestamp: 1734777489810
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
+  sha256: 666c0c431b23c6cec6e492840b176dde533d48b7e6fb8883f5071223433776aa
+  md5: 92ed62436b625154323d40d5f2f11dd7
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - pthread-stubs
+  - xorg-libxau >=1.0.11,<2.0a0
+  - xorg-libxdmcp
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 395888
+  timestamp: 1727278577118
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+  sha256: 6ae68e0b86423ef188196fff6207ed0c8195dd84273cb5623b85aa08033a410c
+  md5: 5aa797f8787fe7a17d1b0821485b5adc
+  depends:
+  - libgcc-ng >=12
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 100393
+  timestamp: 1702724383534
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxkbcommon-1.8.0-hc4a0caf_0.conda
+  sha256: 583203155abcfb03938d8473afbf129156b5b30301a0f796c8ecca8c5b7b2ed2
+  md5: f1656760dbf05f47f962bfdc59fc3416
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - libxcb >=1.17.0,<2.0a0
+  - libxml2 >=2.13.5,<3.0a0
+  - xkeyboard-config
+  - xorg-libxau >=1.0.12,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT/X11 Derivative
+  license_family: MIT
+  purls: []
+  size: 642349
+  timestamp: 1738735301999
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.13.6-h8d12d68_0.conda
+  sha256: db8af71ea9c0ae95b7cb4a0f59319522ed2243942437a1200ceb391493018d85
+  md5: 328382c0e0ca648e5c189d5ec336c604
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - icu >=75.1,<76.0a0
+  - libgcc >=13
+  - libiconv >=1.18,<2.0a0
+  - liblzma >=5.6.4,<6.0a0
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 690296
+  timestamp: 1739952967309
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+  sha256: d4bfe88d7cb447768e31650f06257995601f89076080e76df55e3112d4e47dc4
+  md5: edb0dca6bc32e4f4789199455a1dbeb8
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  constrains:
+  - zlib 1.3.1 *_2
+  arch: x86_64
+  platform: linux
+  license: Zlib
+  license_family: Other
+  purls: []
+  size: 60963
+  timestamp: 1727963148474
+- conda: https://conda.anaconda.org/conda-forge/linux-64/llvm-openmp-19.1.7-h024ca30_0.conda
+  sha256: 5383e32604e03814b6011fa01a5332057934181a7ea0e90abba7890c17cabce6
+  md5: 9915f85a72472011550550623cce2d53
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  constrains:
+  - openmp 19.1.7|19.1.7.*
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0 WITH LLVM-exception
+  license_family: APACHE
+  purls: []
+  size: 3190529
+  timestamp: 1736986301022
+- conda: https://conda.anaconda.org/conda-forge/linux-64/llvmlite-0.43.0-py312h374181b_1.conda
+  sha256: b260285b29834f9b003e2928d778c19b8ed0ca1aff5aa8aa7ec8f21f9b23c2e4
+  md5: ed6ead7e9ab9469629c6cfb363b5c6e2
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libllvm14 >=14.0.6,<14.1.0a0
+  - libstdcxx >=13
+  - libzlib >=1.3.1,<2.0a0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/llvmlite?source=hash-mapping
+  size: 3442782
+  timestamp: 1725305160474
+- conda: https://conda.anaconda.org/conda-forge/noarch/logical-unification-0.4.6-pyhd8ed1ab_0.conda
+  sha256: 2b70aa838779516e05f93158f9f5b15671fc080cec20d05ca0e3a992e391a6e9
+  md5: bd04410bd092c8f62f23a3aea41f47eb
+  depends:
+  - multipledispatch
+  - python >=3.6
+  - toolz
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/logical-unification?source=hash-mapping
+  size: 18160
+  timestamp: 1683416555508
+- conda: https://conda.anaconda.org/conda-forge/linux-64/make-4.4.1-hb9d3cd8_2.conda
+  sha256: d652c7bd4d3b6f82b0f6d063b0d8df6f54cc47531092d7ff008e780f3261bdda
+  md5: 33405d2a66b1411db9f7242c8b97c9e7
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-or-later
+  license_family: GPL
+  purls: []
+  size: 513088
+  timestamp: 1727801714848
+- conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-3.0.0-pyhd8ed1ab_1.conda
+  sha256: 0fbacdfb31e55964152b24d5567e9a9996e1e7902fb08eb7d91b5fd6ce60803a
+  md5: fee3164ac23dfca50cfcc8b85ddefb81
+  depends:
+  - mdurl >=0.1,<1
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/markdown-it-py?source=hash-mapping
+  size: 64430
+  timestamp: 1733250550053
+- conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.0-py312hd3ec401_0.conda
+  sha256: eed67ea988883a3c05160c6d02f34f5a4b6405713cf699d9117eb68fb4743017
+  md5: c27a17a8c54c0d35cf83bbc0de8f7f77
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - contourpy >=1.0.1
+  - cycler >=0.10
+  - fonttools >=4.22.0
+  - freetype >=2.12.1,<3.0a0
+  - kiwisolver >=1.3.1
+  - libgcc >=13
+  - libstdcxx >=13
+  - numpy >=1.19,<3
+  - numpy >=1.23
+  - packaging >=20.0
+  - pillow >=8
+  - pyparsing >=2.3.1
+  - python >=3.12,<3.13.0a0
+  - python-dateutil >=2.7
+  - python_abi 3.12.* *_cp312
+  - qhull >=2020.2,<2020.3.0a0
+  - tk >=8.6.13,<8.7.0a0
+  arch: x86_64
+  platform: linux
+  license: PSF-2.0
+  license_family: PSF
+  purls:
+  - pkg:pypi/matplotlib?source=hash-mapping
+  size: 8210655
+  timestamp: 1734380560683
+- conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.1.7-pyhd8ed1ab_1.conda
+  sha256: 69b7dc7131703d3d60da9b0faa6dd8acbf6f6c396224cf6aef3e855b8c0c41c6
+  md5: af6ab708897df59bd6e7283ceab1b56b
+  depends:
+  - python >=3.9
+  - traitlets
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/matplotlib-inline?source=hash-mapping
+  size: 14467
+  timestamp: 1733417051523
+- conda: https://conda.anaconda.org/conda-forge/linux-64/maturin-1.8.2-py312h6ab59e4_0.conda
+  sha256: 580b95a925375b0aa7cf6ed61501f6464601946b8104fa0bc4d90e3544fadc5a
+  md5: a64484deff2212624f77f4e505e54fbd
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - openssl >=3.4.0,<4.0a0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - tomli >=1.1.0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/maturin?source=hash-mapping
+  size: 6454927
+  timestamp: 1738795289777
+- conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
+  sha256: 78c1bbe1723449c52b7a9df1af2ee5f005209f67e40b6e1d3c7619127c43b1c7
+  md5: 592132998493b3ff25fd7479396e8351
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/mdurl?source=hash-mapping
+  size: 14465
+  timestamp: 1733255681319
+- conda: https://conda.anaconda.org/conda-forge/noarch/minikanren-1.0.3-pyhd8ed1ab_1.conda
+  sha256: 74081492f48c70681cc38e4f85c5f1c1ea5787b25061bcae544b04b5ea7c08c1
+  md5: b48b29582e413aca4de7a366fe98d928
+  depends:
+  - cons >=0.4.0
+  - etuples >=0.3.1
+  - logical-unification >=0.4.1
+  - multipledispatch
+  - python >=3.9
+  - toolz
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/minikanren?source=hash-mapping
+  size: 26338
+  timestamp: 1734538707128
+- conda: https://conda.anaconda.org/conda-forge/linux-64/mkl-2024.2.2-ha957f24_16.conda
+  sha256: 77906b0acead8f86b489da46f53916e624897338770dbf70b04b8f673c9273c1
+  md5: 1459379c79dda834673426504d52b319
+  depends:
+  - _openmp_mutex * *_llvm
+  - _openmp_mutex >=4.5
+  - llvm-openmp >=19.1.2
+  - tbb 2021.*
+  arch: x86_64
+  platform: linux
+  license: LicenseRef-IntelSimplifiedSoftwareOct2022
+  license_family: Proprietary
+  purls: []
+  size: 124718448
+  timestamp: 1730231808335
+- conda: https://conda.anaconda.org/conda-forge/linux-64/mkl-service-2.4.2-py312hf224ee7_0.conda
+  sha256: 405d9f30aca24feb034c0474899613dade19adb52be39efd164d99030c9d3b84
+  md5: 7874d7976d44512eaccba994989017e3
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - mkl >=2024.2.2,<2025.0a0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/mkl-service?source=hash-mapping
+  size: 73420
+  timestamp: 1728843500808
+- pypi: https://files.pythonhosted.org/packages/08/57/5d58fad4124192b1be42f68bd0c0ddaa26e44a730ff8c9337adade2f5632/ml_dtypes-0.5.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+  name: ml-dtypes
+  version: 0.5.1
+  sha256: ad4953c5eb9c25a56d11a913c2011d7e580a435ef5145f804d98efa14477d390
+  requires_dist:
+  - numpy>=1.21
+  - numpy>=1.21.2 ; python_full_version >= '3.10'
+  - numpy>=1.23.3 ; python_full_version >= '3.11'
+  - numpy>=1.26.0 ; python_full_version >= '3.12'
+  - numpy>=2.1.0 ; python_full_version >= '3.13'
+  - absl-py ; extra == 'dev'
+  - pytest ; extra == 'dev'
+  - pytest-xdist ; extra == 'dev'
+  - pylint>=2.6.0 ; extra == 'dev'
+  - pyink ; extra == 'dev'
+  requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/msgpack-python-1.1.0-py312h68727a3_0.conda
+  sha256: 4bc53333774dea1330643b7e23aa34fd6880275737fc2e07491795872d3af8dd
+  md5: 5c9b020a3f86799cdc6115e55df06146
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/msgpack?source=hash-mapping
+  size: 105271
+  timestamp: 1725975182669
+- conda: https://conda.anaconda.org/conda-forge/noarch/multipledispatch-0.6.0-pyhd8ed1ab_1.conda
+  sha256: c6216a21154373b340c64f321f22fec51db4ee6156c2e642fa58368103ac5d09
+  md5: 121a57fce7fff0857ec70fa03200962f
+  depends:
+  - python >=3.6
+  - six
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/multipledispatch?source=hash-mapping
+  size: 17254
+  timestamp: 1721907640382
+- conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyh9f0ad1d_0.tar.bz2
+  sha256: f86fb22b58e93d04b6f25e0d811b56797689d598788b59dcb47f59045b568306
+  md5: 2ba8498c1018c1e9c61eb99b973dfe19
+  depends:
+  - python
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/munkres?source=hash-mapping
+  size: 12452
+  timestamp: 1600387789153
+- conda: https://conda.anaconda.org/conda-forge/noarch/nbclient-0.10.2-pyhd8ed1ab_0.conda
+  sha256: a20cff739d66c2f89f413e4ba4c6f6b59c50d5c30b5f0d840c13e8c9c2df9135
+  md5: 6bb0d77277061742744176ab555b723c
+  depends:
+  - jupyter_client >=6.1.12
+  - jupyter_core >=4.12,!=5.0.*
+  - nbformat >=5.1
+  - python >=3.8
+  - traitlets >=5.4
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/nbclient?source=hash-mapping
+  size: 28045
+  timestamp: 1734628936013
+- conda: https://conda.anaconda.org/conda-forge/noarch/nbformat-5.10.4-pyhd8ed1ab_1.conda
+  sha256: 7a5bd30a2e7ddd7b85031a5e2e14f290898098dc85bea5b3a5bf147c25122838
+  md5: bbe1963f1e47f594070ffe87cdf612ea
+  depends:
+  - jsonschema >=2.6
+  - jupyter_core >=4.12,!=5.0.*
+  - python >=3.9
+  - python-fastjsonschema >=2.15
+  - traitlets >=5.1
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/nbformat?source=hash-mapping
+  size: 100945
+  timestamp: 1733402844974
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+  sha256: 3fde293232fa3fca98635e1167de6b7c7fda83caf24b9d6c91ec9eefb4f4d586
+  md5: 47e340acb35de30501a76c7c799c41d7
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: X11 AND BSD-3-Clause
+  purls: []
+  size: 891641
+  timestamp: 1738195959188
+- conda: https://conda.anaconda.org/conda-forge/noarch/nest-asyncio-1.6.0-pyhd8ed1ab_1.conda
+  sha256: bb7b21d7fd0445ddc0631f64e66d91a179de4ba920b8381f29b9d006a42788c0
+  md5: 598fd7d4d0de2455fb74f56063969a97
+  depends:
+  - python >=3.9
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/nest-asyncio?source=hash-mapping
+  size: 11543
+  timestamp: 1733325673691
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numba-0.60.0-py312h83e6fd3_0.conda
+  sha256: af31c1989ddf1cd46f073f32a8150274c606fdc9fced0e4f5aaf0571b97bd09f
+  md5: e064ca33edf91ac117236c4b5dee207a
+  depends:
+  - _openmp_mutex >=4.5
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  - llvmlite >=0.43.0,<0.44.0a0
+  - numpy >=1.19,<3
+  - numpy >=1.22.3,<2.1
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  constrains:
+  - cuda-version >=11.2
+  - tbb >=2021.6.0
+  - cuda-python >=11.6
+  - scipy >=1.0
+  - libopenblas !=0.3.6
+  - cudatoolkit >=11.2
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/numba?source=hash-mapping
+  size: 5695278
+  timestamp: 1718888170104
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numcodecs-0.15.1-py312hf9745cd_0.conda
+  sha256: 209a84599e36db68865dce5618c3328a2d57267d339255204815885b220a20f2
+  md5: 8a1f88d4985ee1c16b0db1af39a8554d
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - deprecated
+  - libgcc >=13
+  - libstdcxx >=13
+  - msgpack-python
+  - numpy >=1.19,<3
+  - numpy >=1.24
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/numcodecs?source=hash-mapping
+  size: 848654
+  timestamp: 1739285119780
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py312heda63a1_0.conda
+  sha256: fe3459c75cf84dcef6ef14efcc4adb0ade66038ddd27cadb894f34f4797687d8
+  md5: d8285bea2a350f63fab23bf460221f3f
+  depends:
+  - libblas >=3.9.0,<4.0a0
+  - libcblas >=3.9.0,<4.0a0
+  - libgcc-ng >=12
+  - liblapack >=3.9.0,<4.0a0
+  - libstdcxx-ng >=12
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  constrains:
+  - numpy-base <0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/numpy?source=hash-mapping
+  size: 7484186
+  timestamp: 1707225809722
+- conda: https://conda.anaconda.org/conda-forge/noarch/objprint-0.3.0-pyhd8ed1ab_0.conda
+  sha256: ff58f788e9e8c74a6eb2f194b4c18e5bc39a0000172f0b1ec016afae637961f2
+  md5: 8f8399ecb94bd96e0d73e02053525808
+  depends:
+  - python >=3.9
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/objprint?source=hash-mapping
+  size: 38864
+  timestamp: 1731340445369
+- conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.3-h5fbd93e_0.conda
+  sha256: 5bee706ea5ba453ed7fd9da7da8380dd88b865c8d30b5aaec14d2b6dd32dbc39
+  md5: 9e5816bc95d285c115a3ebc2f8563564
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libpng >=1.6.44,<1.7.0a0
+  - libstdcxx >=13
+  - libtiff >=4.7.0,<4.8.0a0
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls: []
+  size: 342988
+  timestamp: 1733816638720
+- conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.4.1-h7b32b05_0.conda
+  sha256: cbf62df3c79a5c2d113247ddea5658e9ff3697b6e741c210656e239ecaf1768f
+  md5: 41adf927e746dc75ecf0ef841c454e48
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - ca-certificates
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls: []
+  size: 2939306
+  timestamp: 1739301879343
+- pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
+  name: opt-einsum
+  version: 3.4.0
+  sha256: 69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd
+  requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/5c/24/28d0bb21600a78e46754947333ec9a297044af884d360092eb8561575fe9/optax-0.2.4-py3-none-any.whl
+  name: optax
+  version: 0.2.4
+  sha256: db35c04e50b52596662efb002334de08c2a0a74971e4da33f467e84fac08886a
+  requires_dist:
+  - absl-py>=0.7.1
+  - chex>=0.1.87
+  - jax>=0.4.27
+  - jaxlib>=0.4.27
+  - numpy>=1.18.0
+  - etils[epy]
+  - sphinx>=6.0.0 ; extra == 'docs'
+  - sphinx-book-theme>=1.0.1 ; extra == 'docs'
+  - sphinxcontrib-katex ; extra == 'docs'
+  - sphinx-autodoc-typehints ; extra == 'docs'
+  - ipython>=8.8.0 ; extra == 'docs'
+  - myst-nb>=1.0.0 ; extra == 'docs'
+  - matplotlib>=3.5.0 ; extra == 'docs'
+  - sphinx-gallery>=0.14.0 ; extra == 'docs'
+  - sphinx-collections>=0.0.1 ; extra == 'docs'
+  - tensorflow>=2.4.0 ; extra == 'docs'
+  - tensorflow-datasets>=4.2.0 ; extra == 'docs'
+  - flax ; extra == 'docs'
+  - sphinx-contributors ; extra == 'docs'
+  - absl-py>=1.0.0 ; extra == 'dp-accounting'
+  - attrs>=21.4.0 ; extra == 'dp-accounting'
+  - mpmath>=1.2.1 ; extra == 'dp-accounting'
+  - numpy>=1.21.4 ; extra == 'dp-accounting'
+  - scipy>=1.7.1 ; extra == 'dp-accounting'
+  - tensorflow-datasets>=4.2.0 ; extra == 'examples'
+  - tensorflow>=2.4.0 ; extra == 'examples'
+  - dp-accounting>=0.4 ; extra == 'examples'
+  - ipywidgets ; extra == 'examples'
+  - flax ; extra == 'examples'
+  - dm-tree>=0.1.7 ; extra == 'test'
+  - flax>=0.5.3 ; extra == 'test'
+  - scipy>=1.7.1 ; extra == 'test'
+  - scikit-learn ; extra == 'test'
+  requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.10.15-py312h12e396e_0.conda
+  sha256: a28431bfb90c99d577bfcd848e05d819e7efe36c50f0b45ed0545f606eb3ee7d
+  md5: 543c7130b6949b6df40e378f9d1db91e
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/orjson?source=hash-mapping
+  size: 303663
+  timestamp: 1737229215301
+- conda: https://conda.anaconda.org/conda-forge/noarch/packaging-24.2-pyhd8ed1ab_2.conda
+  sha256: da157b19bcd398b9804c5c52fc000fcb8ab0525bdb9c70f95beaa0bb42f85af1
+  md5: 3bfed7e6228ebf2f7b9eaa47f1b4e2aa
+  depends:
+  - python >=3.8
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/packaging?source=hash-mapping
+  size: 60164
+  timestamp: 1733203368787
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-2.2.3-py312hf9745cd_1.conda
+  sha256: ad275a83bfebfa8a8fee9b0569aaf6f513ada6a246b2f5d5b85903d8ca61887e
+  md5: 8bce4f6caaf8c5448c7ac86d87e26b4b
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - numpy >=1.19,<3
+  - numpy >=1.22.4
+  - python >=3.12,<3.13.0a0
+  - python-dateutil >=2.8.1
+  - python-tzdata >=2022a
+  - python_abi 3.12.* *_cp312
+  - pytz >=2020.1,<2024.2
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/pandas?source=hash-mapping
+  size: 15436913
+  timestamp: 1726879054912
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pandoc-3.4-ha770c72_0.conda
+  sha256: 8b6c7ddd9422cc6b7ddc10aa8d184f07c1534429e106c441f679f21e95db31c8
+  md5: 61c94057aaa5ae6145137ce1fddb2c04
+  arch: x86_64
+  platform: linux
+  license: GPL-2.0-or-later
+  license_family: GPL
+  purls: []
+  size: 21006632
+  timestamp: 1726013132144
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pango-1.56.1-h861ebed_0.conda
+  sha256: 20e5e280859a7803e8b5a09f18a7e43b56d1b8e61e4888c1a24cbb0d5b9cabd3
+  md5: 59e660508a4de9401543303d5f576aeb
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - cairo >=1.18.2,<2.0a0
+  - fontconfig >=2.15.0,<3.0a0
+  - fonts-conda-ecosystem
+  - freetype >=2.12.1,<3.0a0
+  - fribidi >=1.0.10,<2.0a0
+  - harfbuzz >=10.2.0,<11.0a0
+  - libexpat >=2.6.4,<3.0a0
+  - libgcc >=13
+  - libglib >=2.82.2,<3.0a0
+  - libpng >=1.6.45,<1.7.0a0
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: LGPL-2.1-or-later
+  purls: []
+  size: 451406
+  timestamp: 1737510786003
+- pypi: https://files.pythonhosted.org/packages/20/f8/7b1f5b84f84e43face7349ac2df23f0f75d36b5b4a1c1c0305edce82bcc8/paramax-0.0.0-py3-none-any.whl
+  name: paramax
+  version: 0.0.0
+  sha256: 02d0120e626de300680a1661b138feeba14b418031d6e976fd679db5fd03509a
+  requires_dist:
+  - equinox
+  - jax
+  - jaxtyping
+  - beartype ; extra == 'dev'
+  - pytest ; extra == 'dev'
+  - ruff ; extra == 'dev'
+  - sphinx ; extra == 'dev'
+  - sphinx-autodoc-typehints ; extra == 'dev'
+  - sphinx-book-theme ; extra == 'dev'
+  - sphinx-copybutton ; extra == 'dev'
+  requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.4-pyhd8ed1ab_1.conda
+  sha256: 17131120c10401a99205fc6fe436e7903c0fa092f1b3e80452927ab377239bcc
+  md5: 5c092057b6badd30f75b06244ecd01c9
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/parso?source=hash-mapping
+  size: 75295
+  timestamp: 1733271352153
+- conda: https://conda.anaconda.org/conda-forge/noarch/patsy-1.0.1-pyhd8ed1ab_1.conda
+  sha256: ab52916f056b435757d46d4ce0a93fd73af47df9c11fd72b74cc4b7e1caca563
+  md5: ee23fabfd0a8c6b8d6f3729b47b2859d
+  depends:
+  - numpy >=1.4.0
+  - python >=3.9
+  license: BSD-2-Clause AND PSF-2.0
+  license_family: BSD
+  purls:
+  - pkg:pypi/patsy?source=hash-mapping
+  size: 186594
+  timestamp: 1733792482894
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.44-hba22ea6_2.conda
+  sha256: 1087716b399dab91cc9511d6499036ccdc53eb29a288bebcb19cf465c51d7c0d
+  md5: df359c09c41cd186fffb93a2d87aa6f5
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - bzip2 >=1.0.8,<2.0a0
+  - libgcc-ng >=12
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 952308
+  timestamp: 1723488734144
+- conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
+  sha256: 202af1de83b585d36445dc1fda94266697341994d1a3328fabde4989e1b3d07a
+  md5: d0d408b1f18883a944376da5cf8101ea
+  depends:
+  - ptyprocess >=0.5
+  - python >=3.9
+  license: ISC
+  purls:
+  - pkg:pypi/pexpect?source=compressed-mapping
+  size: 53561
+  timestamp: 1733302019362
+- conda: https://conda.anaconda.org/conda-forge/noarch/pickleshare-0.7.5-pyhd8ed1ab_1004.conda
+  sha256: e2ac3d66c367dada209fc6da43e645672364b9fd5f9d28b9f016e24b81af475b
+  md5: 11a9d1d09a3615fc07c3faf79bc0b943
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pickleshare?source=hash-mapping
+  size: 11748
+  timestamp: 1733327448200
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-11.1.0-py312h80c1187_0.conda
+  sha256: 5c347962202b55ae4d8a463e0555c5c6ca33396266a08284bf1384399894e541
+  md5: d3894405f05b2c0f351d5de3ae26fa9c
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - freetype >=2.12.1,<3.0a0
+  - lcms2 >=2.16,<3.0a0
+  - libgcc >=13
+  - libjpeg-turbo >=3.0.0,<4.0a0
+  - libtiff >=4.7.0,<4.8.0a0
+  - libwebp-base >=1.5.0,<2.0a0
+  - libxcb >=1.17.0,<2.0a0
+  - libzlib >=1.3.1,<2.0a0
+  - openjpeg >=2.5.3,<3.0a0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - tk >=8.6.13,<8.7.0a0
+  arch: x86_64
+  platform: linux
+  license: HPND
+  purls:
+  - pkg:pypi/pillow?source=hash-mapping
+  size: 42749785
+  timestamp: 1735929845390
+- conda: https://conda.anaconda.org/conda-forge/noarch/pip-24.3.1-pyh8b19718_2.conda
+  sha256: da8c8888de10c1e4234ebcaa1550ac2b4b5408ac20f093fe641e4bc8c9c9f3eb
+  md5: 04e691b9fadd93a8a9fad87a81d4fd8f
+  depends:
+  - python >=3.9,<3.13.0a0
+  - setuptools
+  - wheel
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pip?source=hash-mapping
+  size: 1245116
+  timestamp: 1734466348103
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pixman-0.44.2-h29eaf8c_0.conda
+  sha256: 747c58db800d5583fee78e76240bf89cbaeedf7ab1ef339c2990602332b9c4be
+  md5: 5e2a7acfa2c24188af39e7944e1b3604
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 381072
+  timestamp: 1733698987122
+- conda: https://conda.anaconda.org/conda-forge/noarch/pkgutil-resolve-name-1.3.10-pyhd8ed1ab_2.conda
+  sha256: adb2dde5b4f7da70ae81309cce6188ed3286ff280355cf1931b45d91164d2ad8
+  md5: 5a5870a74432aa332f7d32180633ad05
+  depends:
+  - python >=3.9
+  license: MIT AND PSF-2.0
+  purls:
+  - pkg:pypi/pkgutil-resolve-name?source=hash-mapping
+  size: 10693
+  timestamp: 1733344619659
+- conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.3.6-pyhd8ed1ab_1.conda
+  sha256: bb50f6499e8bc1d1a26f17716c97984671121608dc0c3ecd34858112bce59a27
+  md5: 577852c7e53901ddccc7e6a9959ddebe
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/platformdirs?source=hash-mapping
+  size: 20448
+  timestamp: 1733232756001
+- conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.5.0-pyhd8ed1ab_1.conda
+  sha256: 122433fc5318816b8c69283aaf267c73d87aa2d09ce39f64c9805c9a3b264819
+  md5: e9dcbce5f45f9ee500e728ae58b605b6
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pluggy?source=hash-mapping
+  size: 23595
+  timestamp: 1733222855563
+- conda: https://conda.anaconda.org/conda-forge/linux-64/polars-1.22.0-py312hda0fa55_0.conda
+  sha256: 5c9bd84bc4a2fed0e0f59938901e5e36a3da697df296d4301fd9da6a12d984fd
+  md5: ae768211e65e308125e783771938ab5e
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - numpy >=1.16.0
+  - packaging
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/polars?source=hash-mapping
+  size: 24683878
+  timestamp: 1739040110404
+- pypi: https://files.pythonhosted.org/packages/8f/4d/b72e0782abec07f3d8dabf24cf12673d26b173af2046eb4e67365c776ccf/posteriordb-0.2.0-py3-none-any.whl
+  name: posteriordb
+  version: 0.2.0
+  sha256: b6d6f3a349d34db6d4a68da899c818a95e5824c5e23824fc0ebe422f4bd6bac1
+  requires_dist:
+  - requests
+- conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.50-pyha770c72_0.conda
+  sha256: 0749c49a349bf55b8539ce5addce559b77592165da622944a51c630e94d97889
+  md5: 7d823138f550b14ecae927a5ff3286de
+  depends:
+  - python >=3.9
+  - wcwidth
+  constrains:
+  - prompt_toolkit 3.0.50
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/prompt-toolkit?source=hash-mapping
+  size: 271905
+  timestamp: 1737453457168
+- conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-6.1.1-py312h66e93f0_0.conda
+  sha256: 55d4fd0b294aeada0d7810fcc25503b59ec34c4390630789bd61c085b9ce649f
+  md5: add2c79595fa8a9b6d653d7e4e2cf05f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/psutil?source=hash-mapping
+  size: 487053
+  timestamp: 1735327468212
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
+  sha256: 9c88f8c64590e9567c6c80823f0328e58d3b1efb0e1c539c0315ceca764e0973
+  md5: b3c17d95b5a10c6e64a21fa17573e70e
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 8252
+  timestamp: 1726802366959
+- conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
+  sha256: a7713dfe30faf17508ec359e0bc7e0983f5d94682492469bd462cdaae9c64d83
+  md5: 7d9daffbb8d8e0af0f769dbbcd173a54
+  depends:
+  - python >=3.9
+  license: ISC
+  purls:
+  - pkg:pypi/ptyprocess?source=hash-mapping
+  size: 19457
+  timestamp: 1733302371990
+- conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
+  sha256: 71bd24600d14bb171a6321d523486f6a06f855e75e547fa0cb2a0953b02047f0
+  md5: 3bfdfb8dbcdc4af1ae3f9a8eb3948f04
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pure-eval?source=hash-mapping
+  size: 16668
+  timestamp: 1733569518868
+- conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.19.1-pyhd8ed1ab_0.conda
+  sha256: 28a3e3161390a9d23bc02b4419448f8d27679d9e2c250e29849e37749c8de86b
+  md5: 232fb4577b6687b2d503ef8e254270c9
+  depends:
+  - python >=3.9
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/pygments?source=hash-mapping
+  size: 888600
+  timestamp: 1736243563082
+- conda: https://conda.anaconda.org/conda-forge/noarch/pymc-5.20.1-hd8ed1ab_0.conda
+  noarch: python
+  sha256: 2077549cd237ff61273da575c81d289ba4373748688351d688b6dff916c26e59
+  md5: 4bb94e034b8e9408ec4b23a0a11ddb38
+  depends:
+  - pymc-base 5.20.1 pyhd8ed1ab_0
+  - pytensor
+  - python-graphviz
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/pymc?source=compressed-mapping
+  size: 11767
+  timestamp: 1738867855599
+- conda: https://conda.anaconda.org/conda-forge/noarch/pymc-base-5.20.1-pyhd8ed1ab_0.conda
+  sha256: 44009910a448f322bdd3d906324f4b7e5aa3a0865cabdc0a66c5b3e254363c1f
+  md5: dfe6d235c12d3bd0145f3aed3976c1dc
+  depends:
+  - arviz >=0.13.0
+  - cachetools >=4.2.1
+  - cloudpickle
+  - numpy >=1.25.0
+  - pandas >=0.24.0
+  - pytensor-base >=2.26.1,<2.28
+  - python >=3.10
+  - rich >=13.7.1
+  - scipy >=1.4.1
+  - threadpoolctl >=3.1.0,<4.0.0
+  - typing-extensions >=3.7.4
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/pymc?source=hash-mapping
+  size: 347711
+  timestamp: 1738867852390
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.2.1-pyhd8ed1ab_0.conda
+  sha256: f513fed4001fd228d3bf386269237b4ca6bff732c99ffc11fcbad8529b35407c
+  md5: 285e237b8f351e85e7574a2c7bfa6d46
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pyparsing?source=hash-mapping
+  size: 93082
+  timestamp: 1735698406955
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pytensor-2.27.1-py312h97902ae_0.conda
+  sha256: b52c0e3076fba01e2001d70237d85b698525a22bfaee13a15c02cd375b0a08e0
+  md5: 3dcb98f87392898b974cfca9d581e91c
+  depends:
+  - blas
+  - gcc_linux-64 13.*
+  - gxx
+  - gxx_linux-64 13.*
+  - mkl-service
+  - pytensor-base 2.27.1 py312h25a0e75_0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - sysroot_linux-64 2.17.*
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 9130
+  timestamp: 1738648470447
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pytensor-base-2.27.1-py312h25a0e75_0.conda
+  sha256: 96b752c66a1b83c56277df8446a5f0692a2d6dff97e1b929ed257d1249f5c414
+  md5: e80d864b4d2259dbdd420783aff9b97b
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - cons
+  - etuples
+  - filelock >=3.15
+  - libgcc >=13
+  - libstdcxx >=13
+  - logical-unification
+  - minikanren
+  - numpy >=1.26.4,<2.0a0
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - scipy >=1,<2
+  - setuptools >=59.0.0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/pytensor?source=hash-mapping
+  size: 2321283
+  timestamp: 1738648453740
+- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-8.3.4-pyhd8ed1ab_1.conda
+  sha256: 75245ca9d0cbd6d38bb45ec02430189a9d4c21c055c5259739d738a2298d61b3
+  md5: 799ed216dc6af62520f32aa39bc1c2bb
+  depends:
+  - colorama
+  - exceptiongroup >=1.0.0rc8
+  - iniconfig
+  - packaging
+  - pluggy <2,>=1.5
+  - python >=3.9
+  - tomli >=1
+  constrains:
+  - pytest-faulthandler >=2
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pytest?source=hash-mapping
+  size: 259195
+  timestamp: 1733217599806
+- conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.9-h9e4cc4f_0_cpython.conda
+  sha256: 64fed5178f1e9c8ac0f572ac0ce37955f5dee7b2bcac665202bc14f1f7dd618a
+  md5: 5665f0079432f8848079c811cdb537d5
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - bzip2 >=1.0.8,<2.0a0
+  - ld_impl_linux-64 >=2.36.1
+  - libexpat >=2.6.4,<3.0a0
+  - libffi >=3.4,<4.0a0
+  - libgcc >=13
+  - liblzma >=5.6.4,<6.0a0
+  - libnsl >=2.0.1,<2.1.0a0
+  - libsqlite >=3.48.0,<4.0a0
+  - libuuid >=2.38.1,<3.0a0
+  - libxcrypt >=4.4.36
+  - libzlib >=1.3.1,<2.0a0
+  - ncurses >=6.5,<7.0a0
+  - openssl >=3.4.1,<4.0a0
+  - readline >=8.2,<9.0a0
+  - tk >=8.6.13,<8.7.0a0
+  - tzdata
+  constrains:
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: Python-2.0
+  purls: []
+  size: 31581682
+  timestamp: 1739521496324
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhff2d567_1.conda
+  sha256: a50052536f1ef8516ed11a844f9413661829aa083304dc624c5925298d078d79
+  md5: 5ba79d7c71f03c678c8ead841f347d6e
+  depends:
+  - python >=3.9
+  - six >=1.5
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/python-dateutil?source=hash-mapping
+  size: 222505
+  timestamp: 1733215763718
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.1-pyhd8ed1ab_0.conda
+  sha256: 1b09a28093071c1874862422696429d0d35bd0b8420698003ac004746c5e82a2
+  md5: 38e34d2d1d9dca4fb2b9a0a04f604e2c
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/fastjsonschema?source=hash-mapping
+  size: 226259
+  timestamp: 1733236073335
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-graphviz-0.20.3-pyh91182bf_2.conda
+  sha256: c8f5d3d23b5962524217f33549add8d6c5af22fe839b49603f4588771154a51c
+  md5: f822f0e13849c2283f72ec4aa120eeaa
+  depends:
+  - graphviz >=2.46.1
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/graphviz?source=hash-mapping
+  size: 38220
+  timestamp: 1733792086212
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.1-pyhd8ed1ab_0.conda
+  sha256: 1597d6055d34e709ab8915091973552a0b8764c8032ede07c4e99670da029629
+  md5: 392c91c42edd569a7ec99ed8648f597a
+  depends:
+  - python >=3.9
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/tzdata?source=hash-mapping
+  size: 143794
+  timestamp: 1737541204030
+- conda: https://conda.anaconda.org/conda-forge/linux-64/python_abi-3.12-5_cp312.conda
+  build_number: 5
+  sha256: d10e93d759931ffb6372b45d65ff34d95c6000c61a07e298d162a3bc2accebb0
+  md5: 0424ae29b104430108f5218a66db7260
+  constrains:
+  - python 3.12.* *_cpython
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 6238
+  timestamp: 1723823388266
+- conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2024.1-pyhd8ed1ab_0.conda
+  sha256: 1a7d6b233f7e6e3bbcbad054c8fd51e690a67b129a899a056a5e45dd9f00cb41
+  md5: 3eeeeb9e4827ace8c0c1419c85d590ad
+  depends:
+  - python >=3.7
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pytz?source=hash-mapping
+  size: 188538
+  timestamp: 1706886944988
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.2-py312h178313f_2.conda
+  sha256: 159cba13a93b3fe084a1eb9bda0a07afc9148147647f0d437c3c3da60980503b
+  md5: cf2485f39740de96e2a7f2bb18ed2fee
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - yaml >=0.2.5,<0.3.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/pyyaml?source=hash-mapping
+  size: 206903
+  timestamp: 1737454910324
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-26.2.1-py312hbf22597_0.conda
+  sha256: 90ec0da0317d3d76990a40c61e1709ef859dd3d8c63838bad2814f46a63c8a2e
+  md5: 7cec8d0dac15a2d9fea8e49879aa779d
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libsodium >=1.0.20,<1.0.21.0a0
+  - libstdcxx >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - zeromq >=4.3.5,<4.4.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/pyzmq?source=hash-mapping
+  size: 382698
+  timestamp: 1738271121975
+- conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
+  sha256: 776363493bad83308ba30bcb88c2552632581b143e8ee25b1982c8c743e73abc
+  md5: 353823361b1d27eb3960efb076dfcaf6
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc-ng >=12
+  - libstdcxx-ng >=12
+  arch: x86_64
+  platform: linux
+  license: LicenseRef-Qhull
+  purls: []
+  size: 552937
+  timestamp: 1720813982144
+- conda: https://conda.anaconda.org/conda-forge/linux-64/quarto-1.6.40-ha770c72_0.conda
+  sha256: 7d1538903ae2cde212eb2075b2b745ae0c8e5386bae458aebb67311f4d63c126
+  md5: 9b16208e735b988727af946e049af3e3
+  depends:
+  - dart-sass
+  - deno >=1.46.3,<1.46.4.0a0
+  - deno-dom >=0.1.41,<0.1.42.0a0
+  - esbuild
+  - pandoc 3.4
+  - typst 0.11.0
+  arch: x86_64
+  platform: linux
+  license: GPL-2.0-or-later
+  license_family: GPL
+  purls: []
+  size: 15989616
+  timestamp: 1736717196147
+- conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.2-h8228510_1.conda
+  sha256: 5435cf39d039387fbdc977b0a762357ea909a7694d9528ab40f005e9208744d7
+  md5: 47d31b792659ce70f470b5c82fdfb7a4
+  depends:
+  - libgcc-ng >=12
+  - ncurses >=6.3,<7.0a0
+  arch: x86_64
+  platform: linux
+  license: GPL-3.0-only
+  license_family: GPL
+  purls: []
+  size: 281456
+  timestamp: 1679532220005
+- conda: https://conda.anaconda.org/conda-forge/noarch/referencing-0.36.2-pyh29332c3_0.conda
+  sha256: e20909f474a6cece176dfc0dc1addac265deb5fa92ea90e975fbca48085b20c3
+  md5: 9140f1c09dd5489549c6a33931b943c7
+  depends:
+  - attrs >=22.2.0
+  - python >=3.9
+  - rpds-py >=0.7.0
+  - typing_extensions >=4.4.0
+  - python
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/referencing?source=hash-mapping
+  size: 51668
+  timestamp: 1737836872415
+- pypi: https://files.pythonhosted.org/packages/f9/9b/335f9764261e915ed497fcdeb11df5dfd6f7bf257d4a6a2a686d80da4d54/requests-2.32.3-py3-none-any.whl
+  name: requests
+  version: 2.32.3
+  sha256: 70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
+  requires_dist:
+  - charset-normalizer>=2,<4
+  - idna>=2.5,<4
+  - urllib3>=1.21.1,<3
+  - certifi>=2017.4.17
+  - pysocks>=1.5.6,!=1.5.7 ; extra == 'socks'
+  - chardet>=3.0.2,<6 ; extra == 'use-chardet-on-py3'
+  requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/rich-13.9.4-pyhd8ed1ab_1.conda
+  sha256: 06a760c5ae572e72e865d5a87e9fe3cc171e1a9c996e63daf3db52ff1a0b4457
+  md5: 7aed65d4ff222bfb7335997aa40b7da5
+  depends:
+  - markdown-it-py >=2.2.0
+  - pygments >=2.13.0,<3.0.0
+  - python >=3.9
+  - typing_extensions >=4.0.0,<5.0.0
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/rich?source=hash-mapping
+  size: 185646
+  timestamp: 1733342347277
+- conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.22.3-py312h12e396e_0.conda
+  sha256: e8662d21ca3c912ac8941725392b838a29458b106ef22d9489cdf0f8de145fad
+  md5: bfb49da0cc9098597d527def04d66f8b
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  constrains:
+  - __glibc >=2.17
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/rpds-py?source=hash-mapping
+  size: 354410
+  timestamp: 1733366814237
+- pypi: git+https://github.com/aseyboldt/samplerlab#9edea2da1c5b336e53b9238440a2317cf1baba14
+  name: samplerlab
+  version: '0.1'
+  requires_dist:
+  - numpy
+  - arviz
+  - pymc
+  requires_python: '>=3.11'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.15.2-py312ha707e6e_0.conda
+  sha256: b9faaa024b77a3678a988c5a490f02c4029c0d5903998b585100e05bc7d4ff36
+  md5: 00b999c5f9d01fb633db819d79186bd4
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libblas >=3.9.0,<4.0a0
+  - libcblas >=3.9.0,<4.0a0
+  - libgcc >=13
+  - libgfortran
+  - libgfortran5 >=13.3.0
+  - liblapack >=3.9.0,<4.0a0
+  - libstdcxx >=13
+  - numpy <2.5
+  - numpy >=1.19,<3
+  - numpy >=1.23.5
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/scipy?source=hash-mapping
+  size: 17064784
+  timestamp: 1739791925628
+- conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-0.13.2-hd8ed1ab_3.conda
+  noarch: python
+  sha256: ea29a69b14dd6be5cdeeaa551bf50d78cafeaf0351e271e358f9b820fcab4cb0
+  md5: 62afb877ca2c2b4b6f9ecb37320085b6
+  depends:
+  - seaborn-base 0.13.2 pyhd8ed1ab_3
+  - statsmodels >=0.12
+  license: BSD-3-Clause
+  license_family: BSD
+  purls: []
+  size: 6876
+  timestamp: 1733730113224
+- conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-base-0.13.2-pyhd8ed1ab_3.conda
+  sha256: f209c9c18187570b85ec06283c72d64b8738f825b1b82178f194f4866877f8aa
+  md5: fd96da444e81f9e6fcaac38590f3dd42
+  depends:
+  - matplotlib-base >=3.4,!=3.6.1
+  - numpy >=1.20,!=1.24.0
+  - pandas >=1.2
+  - python >=3.9
+  - scipy >=1.7
+  constrains:
+  - seaborn =0.13.2=*_3
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/seaborn?source=hash-mapping
+  size: 227843
+  timestamp: 1733730112409
+- conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-75.8.0-pyhff2d567_0.conda
+  sha256: e0778e4f276e9a81b51c56f51ec22a27b4d8fc955abc0be77ad09ca9bea06bb9
+  md5: 8f28e299c11afdd79e0ec1e279dcdc52
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/setuptools?source=hash-mapping
+  size: 775598
+  timestamp: 1736512753595
+- conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhd8ed1ab_0.conda
+  sha256: 41db0180680cc67c3fa76544ffd48d6a5679d96f4b71d7498a759e94edc9a2db
+  md5: a451d576819089b0d672f18768be0f65
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/six?source=hash-mapping
+  size: 16385
+  timestamp: 1733381032766
+- conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
+  sha256: 570da295d421661af487f1595045760526964f41471021056e993e73089e9c41
+  md5: b1b505328da7a6b246787df4b5a49fbc
+  depends:
+  - asttokens
+  - executing
+  - pure_eval
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/stack-data?source=hash-mapping
+  size: 26988
+  timestamp: 1733569565672
+- conda: https://conda.anaconda.org/conda-forge/noarch/stanio-0.5.1-pyhd8ed1ab_1.conda
+  sha256: da73aba3babd557592d62f5c26646e16c386e4e8077466c818780f4b3060795a
+  md5: a4b46bbc8e0159f7970afef0ca758fdd
+  depends:
+  - numpy
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/stanio?source=hash-mapping
+  size: 13565
+  timestamp: 1734581732441
+- conda: https://conda.anaconda.org/conda-forge/linux-64/statsmodels-0.14.4-py312hc0a28a1_0.conda
+  sha256: 6cc65ba902b32207e8a697b0e0408a28d6cc166be04f1882c40739a86a253d22
+  md5: 97dc960f3d9911964d73c2cf240baea5
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - numpy <3,>=1.22.3
+  - numpy >=1.19,<3
+  - packaging >=21.3
+  - pandas !=2.1.0,>=1.4
+  - patsy >=0.5.6
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  - scipy !=1.9.2,>=1.8
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/statsmodels?source=hash-mapping
+  size: 12103203
+  timestamp: 1727987129263
+- conda: https://conda.anaconda.org/conda-forge/noarch/sysroot_linux-64-2.17-h0157908_18.conda
+  sha256: 69ab5804bdd2e8e493d5709eebff382a72fab3e9af6adf93a237ccf8f7dbd624
+  md5: 460eba7851277ec1fd80a1a24080787a
+  depends:
+  - kernel-headers_linux-64 3.10.0 he073ed8_18
+  - tzdata
+  license: LGPL-2.0-or-later AND LGPL-2.0-or-later WITH exceptions AND GPL-2.0-or-later AND MPL-2.0
+  license_family: GPL
+  purls: []
+  size: 15166921
+  timestamp: 1735290488259
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tbb-2021.13.0-hceb3a55_1.conda
+  sha256: 65463732129899770d54b1fbf30e1bb82fdebda9d7553caf08d23db4590cd691
+  md5: ba7726b8df7b9d34ea80e82b097a4893
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libhwloc >=2.11.2,<2.11.3.0a0
+  - libstdcxx >=13
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: APACHE
+  purls: []
+  size: 175954
+  timestamp: 1732982638805
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tbb-devel-2021.13.0-h1f99690_1.conda
+  sha256: a4f4d9c17057e047406f855a31093183cb4e30bda221506e024cf826e29f9399
+  md5: 0200638907d55c2aadc1ca9d39566316
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - tbb 2021.13.0 hceb3a55_1
+  arch: x86_64
+  platform: linux
+  purls: []
+  size: 1053549
+  timestamp: 1732982652385
+- conda: https://conda.anaconda.org/conda-forge/noarch/threadpoolctl-3.5.0-pyhc1e730c_0.conda
+  sha256: 45e402941f6bed094022c5726a2ca494e6224b85180d2367fb6ddd9aea68079d
+  md5: df68d78237980a159bd7149f33c0e8fd
+  depends:
+  - python >=3.8
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/threadpoolctl?source=hash-mapping
+  size: 23548
+  timestamp: 1714400228771
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h4845f30_101.conda
+  sha256: e0569c9caa68bf476bead1bed3d79650bb080b532c64a4af7d8ca286c08dea4e
+  md5: d453b98d9c83e71da0741bb0ff4d76bc
+  depends:
+  - libgcc-ng >=12
+  - libzlib >=1.2.13,<2.0.0a0
+  arch: x86_64
+  platform: linux
+  license: TCL
+  license_family: BSD
+  purls: []
+  size: 3318875
+  timestamp: 1699202167581
+- conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.2.1-pyhd8ed1ab_1.conda
+  sha256: 18636339a79656962723077df9a56c0ac7b8a864329eb8f847ee3d38495b863e
+  md5: ac944244f1fed2eb49bae07193ae8215
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/tomli?source=hash-mapping
+  size: 19167
+  timestamp: 1733256819729
+- conda: https://conda.anaconda.org/conda-forge/noarch/toolz-1.0.0-pyhd8ed1ab_1.conda
+  sha256: eda38f423c33c2eaeca49ed946a8d3bf466cc3364970e083a65eb2fd85258d87
+  md5: 40d0ed782a8aaa16ef248e68c06c168d
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/toolz?source=hash-mapping
+  size: 52475
+  timestamp: 1733736126261
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.4.2-py312h66e93f0_0.conda
+  sha256: 062a3a3a37fa8615ce57929ba7e982c76f5a5810bcebd435950f6d6c4147c310
+  md5: e417822cb989e80a0d2b1b576fdd1657
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/tornado?source=hash-mapping
+  size: 840414
+  timestamp: 1732616043734
+- conda: https://conda.anaconda.org/conda-forge/noarch/tqdm-4.67.1-pyhd8ed1ab_1.conda
+  sha256: 11e2c85468ae9902d24a27137b6b39b4a78099806e551d390e394a8c34b48e40
+  md5: 9efbfdc37242619130ea42b1cc4ed861
+  depends:
+  - colorama
+  - python >=3.9
+  license: MPL-2.0 or MIT
+  purls:
+  - pkg:pypi/tqdm?source=hash-mapping
+  size: 89498
+  timestamp: 1735661472632
+- conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
+  sha256: f39a5620c6e8e9e98357507262a7869de2ae8cc07da8b7f84e517c9fd6c2b959
+  md5: 019a7385be9af33791c989871317e1ed
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/traitlets?source=hash-mapping
+  size: 110051
+  timestamp: 1733367480074
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.12.2-hd8ed1ab_1.conda
+  noarch: python
+  sha256: c8e9c1c467b5f960b627d7adc1c65fece8e929a3de89967e91ef0f726422fd32
+  md5: b6a408c64b78ec7b779a3e5c7a902433
+  depends:
+  - typing_extensions 4.12.2 pyha770c72_1
+  license: PSF-2.0
+  license_family: PSF
+  purls: []
+  size: 10075
+  timestamp: 1733188758872
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.12.2-pyha770c72_1.conda
+  sha256: 337be7af5af8b2817f115b3b68870208b30c31d3439bec07bfb2d8f4823e3568
+  md5: d17f13df8b65464ca316cbc000a3cb64
+  depends:
+  - python >=3.9
+  license: PSF-2.0
+  license_family: PSF
+  purls:
+  - pkg:pypi/typing-extensions?source=hash-mapping
+  size: 39637
+  timestamp: 1733188758212
+- conda: https://conda.anaconda.org/conda-forge/linux-64/typst-0.11.0-he8a937b_0.conda
+  sha256: e0be4b66c486a15341d7309744ca94c51dbf0c3914c05cbf2bd796aac0c510c5
+  md5: f6ac8c8ab4a4c9ed0ec710e3f5fa3954
+  depends:
+  - libgcc-ng >=12
+  - openssl >=3.2.1,<4.0a0
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls: []
+  size: 11632673
+  timestamp: 1710532616618
+- conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025a-h78e105d_0.conda
+  sha256: c4b1ae8a2931fe9b274c44af29c5475a85b37693999f8c792dad0f8c6734b1de
+  md5: dbcace4706afdfb7eb891f7b37d07c04
+  license: LicenseRef-Public-Domain
+  purls: []
+  size: 122921
+  timestamp: 1737119101255
+- conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-16.0.0-py312h66e93f0_0.conda
+  sha256: 638916105a836973593547ba5cf4891d1f2cb82d1cf14354fcef93fd5b941cdc
+  md5: 617f5d608ff8c28ad546e5d9671cbb95
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/unicodedata2?source=compressed-mapping
+  size: 404401
+  timestamp: 1736692621599
+- pypi: https://files.pythonhosted.org/packages/c8/19/4ec628951a74043532ca2cf5d97b7b14863931476d117c471e8e2b1eb39f/urllib3-2.3.0-py3-none-any.whl
+  name: urllib3
+  version: 2.3.0
+  sha256: 1cee9ad369867bfdbbb48b7dd50374c0967a0bb7710050facf0dd6911440e3df
+  requires_dist:
+  - brotli>=1.0.9 ; platform_python_implementation == 'CPython' and extra == 'brotli'
+  - brotlicffi>=0.8.0 ; platform_python_implementation != 'CPython' and extra == 'brotli'
+  - h2>=4,<5 ; extra == 'h2'
+  - pysocks>=1.5.6,!=1.5.7,<2.0 ; extra == 'socks'
+  - zstandard>=0.18.0 ; extra == 'zstd'
+  requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/viztracer-1.0.2-py312h0f154a2_0.conda
+  sha256: bf1a01de77a6dfc882ae04773b9b457f52116d2fea6ac931328e0b2221319e5d
+  md5: 34b7b85b791656cd249e0ad0c1d7ccb2
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc
+  - libgcc-ng >=12
+  - objprint >=0.3.0
+  - orjson
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/viztracer?source=hash-mapping
+  size: 11356383
+  timestamp: 1738891079447
+- pypi: https://files.pythonhosted.org/packages/39/3b/5b918a0da0d6920e7f7328cf0ab00df31b905d709f458596304f09096785/wadler_lindig-0.1.3-py3-none-any.whl
+  name: wadler-lindig
+  version: 0.1.3
+  sha256: 3018e4e6b115a7ef21c77414a41cbe7e03e83f6b5e25004958e33432a17f3c94
+  requires_dist:
+  - numpy ; extra == 'dev'
+  - pre-commit ; extra == 'dev'
+  - pytest ; extra == 'dev'
+  - jinja2==3.0.3 ; extra == 'docs'
+  - mkdocs-autorefs==1.0.1 ; extra == 'docs'
+  - mkdocs-include-exclude-files==0.0.1 ; extra == 'docs'
+  - mkdocs-material-extensions==1.3.1 ; extra == 'docs'
+  - mkdocs-material==7.3.6 ; extra == 'docs'
+  - mkdocs==1.3.0 ; extra == 'docs'
+  - mkdocstrings==0.17.0 ; extra == 'docs'
+  - mknotebooks==0.7.1 ; extra == 'docs'
+  - nbconvert==6.5.0 ; extra == 'docs'
+  - nbformat==5.4.0 ; extra == 'docs'
+  - pygments==2.14.0 ; extra == 'docs'
+  - pymdown-extensions==9.4 ; extra == 'docs'
+  - pytkdocs-tweaks==0.0.8 ; extra == 'docs'
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/c3/ac/1847959031ce09eb65a70b16d0f07633d7ddaf8a97752cc924f9a8cbe253/watermark-2.5.0-py2.py3-none-any.whl
+  name: watermark
+  version: 2.5.0
+  sha256: d775afa075f5ea15152ad9f4c82dfc95bffb9d0cb4341c9b62073b627191464d
+  requires_dist:
+  - ipython>=6.0
+  - importlib-metadata>=1.4
+  - setuptools
+  - py3nvml>=0.2 ; extra == 'gpu'
+  requires_python: '>=3.7'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/wayland-1.23.1-h3e06ad9_0.conda
+  sha256: 0884b2023a32d2620192cf2e2fc6784b8d1e31cf9f137e49e00802d4daf7d1c1
+  md5: 0a732427643ae5e0486a727927791da1
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libexpat >=2.6.2,<3.0a0
+  - libffi >=3.4,<4.0a0
+  - libgcc-ng >=13
+  - libstdcxx-ng >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 321561
+  timestamp: 1724530461598
+- conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.13-pyhd8ed1ab_1.conda
+  sha256: f21e63e8f7346f9074fd00ca3b079bd3d2fa4d71f1f89d5b6934bf31446dc2a5
+  md5: b68980f2495d096e71c7fd9d7ccf63e6
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/wcwidth?source=hash-mapping
+  size: 32581
+  timestamp: 1733231433877
+- conda: https://conda.anaconda.org/conda-forge/noarch/wheel-0.45.1-pyhd8ed1ab_1.conda
+  sha256: 1b34021e815ff89a4d902d879c3bd2040bc1bd6169b32e9427497fa05c55f1ce
+  md5: 75cb7132eb58d97896e173ef12ac9986
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/wheel?source=hash-mapping
+  size: 62931
+  timestamp: 1733130309598
+- conda: https://conda.anaconda.org/conda-forge/noarch/widgetsnbextension-4.0.13-pyhd8ed1ab_1.conda
+  sha256: a750202ae2a31d8e5ee5a5c127fcc7fa783cd0fbedbc0bf1ab549a109881fa9f
+  md5: 237db148cc37a466e4222d589029b53e
+  depends:
+  - python >=3.9
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/widgetsnbextension?source=hash-mapping
+  size: 898402
+  timestamp: 1733128654300
+- conda: https://conda.anaconda.org/conda-forge/linux-64/wrapt-1.17.2-py312h66e93f0_0.conda
+  sha256: ed3a1700ecc5d38c7e7dc7d2802df1bc1da6ba3d6f6017448b8ded0affb4ae00
+  md5: 669e63af87710f8d52fdec9d4d63b404
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - python >=3.12,<3.13.0a0
+  - python_abi 3.12.* *_cp312
+  arch: x86_64
+  platform: linux
+  license: BSD-2-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/wrapt?source=hash-mapping
+  size: 63590
+  timestamp: 1736869574299
+- conda: https://conda.anaconda.org/conda-forge/noarch/xarray-2025.1.2-pyhd8ed1ab_0.conda
+  sha256: 0f59c2718573770b01d849e05a56a7fe1461f55bf7525c4df5552079c5c03427
+  md5: b8d9af89c48fa3359f05f3324809fcde
+  depends:
+  - numpy >=1.24
+  - packaging >=23.2
+  - pandas >=2.1
+  - python >=3.10
+  constrains:
+  - pint >=0.22
+  - netcdf4 >=1.6.0
+  - cartopy >=0.22
+  - iris >=3.7
+  - h5py >=3.8
+  - nc-time-axis >=1.4
+  - sparse >=0.14
+  - bottleneck >=1.3
+  - scipy >=1.11
+  - numba >=0.57
+  - toolz >=0.12
+  - h5netcdf >=1.3
+  - distributed >=2023.11
+  - dask-core >=2023.11
+  - flox >=0.7
+  - seaborn-base >=0.13
+  - zarr >=2.16
+  - matplotlib-base >=3.8
+  - cftime >=1.6
+  - hdf5 >=1.12
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/xarray?source=hash-mapping
+  size: 837969
+  timestamp: 1738313762187
+- conda: https://conda.anaconda.org/conda-forge/noarch/xarray-einstats-0.8.0-pyhd8ed1ab_1.conda
+  sha256: 8928c391750b3376bf579b5b24db306872618765ebfe17c3fafc3bf09977dab0
+  md5: 469d83261349ad6f3829a9aeebeb8822
+  depends:
+  - numpy >=1.23
+  - python >=3.10
+  - scipy >=1.9
+  - xarray >=2022.09.0
+  license: Apache-2.0
+  license_family: APACHE
+  purls:
+  - pkg:pypi/xarray-einstats?source=hash-mapping
+  size: 33933
+  timestamp: 1734625062168
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xkeyboard-config-2.43-hb9d3cd8_0.conda
+  sha256: 0d89b5873515a1f05d311f37ea4e087bbccc0418afa38f2f6189e97280db3179
+  md5: f725c7425d6d7c15e31f3b99a88ea02f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 389475
+  timestamp: 1727840188958
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libice-1.1.2-hb9d3cd8_0.conda
+  sha256: c12396aabb21244c212e488bbdc4abcdef0b7404b15761d9329f5a4a39113c4b
+  md5: fb901ff28063514abb6046c9ec2c4a45
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 58628
+  timestamp: 1734227592886
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libsm-1.2.5-he73a12e_0.conda
+  sha256: 760f43df6c2ce8cbbbcb8f2f3b7fc0f306716c011e28d1d340f3dfa8ccf29185
+  md5: 4c3e9fab69804ec6077697922d70c6e2
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libuuid >=2.38.1,<3.0a0
+  - xorg-libice >=1.1.2,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 27198
+  timestamp: 1734229639785
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libx11-1.8.11-h4f16b4b_0.conda
+  sha256: a0e7fca9e341dc2455b20cd320fc1655e011f7f5f28367ecf8617cccd4bb2821
+  md5: b6eb6d0cb323179af168df8fe16fb0a1
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libxcb >=1.17.0,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 835157
+  timestamp: 1738613163812
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb9d3cd8_0.conda
+  sha256: ed10c9283974d311855ae08a16dfd7e56241fac632aec3b92e3cfe73cff31038
+  md5: f6ebe2cb3f82ba6c057dde5d9debe4f7
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 14780
+  timestamp: 1734229004433
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcomposite-0.4.6-hb9d3cd8_2.conda
+  sha256: 753f73e990c33366a91fd42cc17a3d19bb9444b9ca5ff983605fa9e953baf57f
+  md5: d3c295b50f092ab525ffe3c2aa4b7413
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxfixes >=6.0.1,<7.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 13603
+  timestamp: 1727884600744
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcursor-1.2.3-hb9d3cd8_0.conda
+  sha256: 832f538ade441b1eee863c8c91af9e69b356cd3e9e1350fff4fe36cc573fc91a
+  md5: 2ccd714aa2242315acaf0a67faea780b
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxfixes >=6.0.1,<7.0a0
+  - xorg-libxrender >=0.9.11,<0.10.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 32533
+  timestamp: 1730908305254
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdamage-1.1.6-hb9d3cd8_0.conda
+  sha256: 43b9772fd6582bf401846642c4635c47a9b0e36ca08116b3ec3df36ab96e0ec0
+  md5: b5fcc7172d22516e1f965490e65e33a4
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxfixes >=6.0.1,<7.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 13217
+  timestamp: 1727891438799
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb9d3cd8_0.conda
+  sha256: 6b250f3e59db07c2514057944a3ea2044d6a8cdde8a47b6497c254520fade1ee
+  md5: 8035c64cb77ed555e3f150b7b3972480
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 19901
+  timestamp: 1727794976192
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxext-1.3.6-hb9d3cd8_0.conda
+  sha256: da5dc921c017c05f38a38bd75245017463104457b63a1ce633ed41f214159c14
+  md5: febbab7d15033c913d53c7a2c102309d
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 50060
+  timestamp: 1727752228921
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxfixes-6.0.1-hb9d3cd8_0.conda
+  sha256: 2fef37e660985794617716eb915865ce157004a4d567ed35ec16514960ae9271
+  md5: 4bdb303603e9821baf5fe5fdff1dc8f8
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 19575
+  timestamp: 1727794961233
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxi-1.8.2-hb9d3cd8_0.conda
+  sha256: 1a724b47d98d7880f26da40e45f01728e7638e6ec69f35a3e11f92acd05f9e7a
+  md5: 17dcc85db3c7886650b8908b183d6876
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxfixes >=6.0.1,<7.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 47179
+  timestamp: 1727799254088
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxinerama-1.1.5-h5888daf_1.conda
+  sha256: 1b9141c027f9d84a9ee5eb642a0c19457c788182a5a73c5a9083860ac5c20a8c
+  md5: 5e2eb9bf77394fc2e5918beefec9f9ab
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 13891
+  timestamp: 1727908521531
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrandr-1.5.4-hb9d3cd8_0.conda
+  sha256: ac0f037e0791a620a69980914a77cb6bb40308e26db11698029d6708f5aa8e0d
+  md5: 2de7f99d6581a4a7adbff607b5c278ca
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxrender >=0.9.11,<0.10.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 29599
+  timestamp: 1727794874300
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrender-0.9.12-hb9d3cd8_0.conda
+  sha256: 044c7b3153c224c6cedd4484dd91b389d2d7fd9c776ad0f4a34f099b3389f4a1
+  md5: 96d57aba173e878a2089d5638016dc5e
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 33005
+  timestamp: 1734229037766
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxtst-1.2.5-hb9d3cd8_3.conda
+  sha256: 752fdaac5d58ed863bbf685bb6f98092fe1a488ea8ebb7ed7b606ccfce08637a
+  md5: 7bbe9a0cc0df0ac5f5a8ad6d6a11af2f
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - xorg-libx11 >=1.8.10,<2.0a0
+  - xorg-libxext >=1.3.6,<2.0a0
+  - xorg-libxi >=1.7.10,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 32808
+  timestamp: 1727964811275
+- conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h7f98852_2.tar.bz2
+  sha256: a4e34c710eeb26945bdbdaba82d3d74f60a78f54a874ec10d373811a5d217535
+  md5: 4cb3ad778ec2d5a7acbdf254eb1c42ae
+  depends:
+  - libgcc-ng >=9.4.0
+  arch: x86_64
+  platform: linux
+  license: MIT
+  license_family: MIT
+  purls: []
+  size: 89141
+  timestamp: 1641346969816
+- conda: https://conda.anaconda.org/conda-forge/noarch/zarr-2.18.4-pyhd8ed1ab_0.conda
+  sha256: a189f686ad211bd20dd303df9f446bdb03b13387be86c4281e66a7830e4f04ff
+  md5: 82725d5fde7cc349b62e5f0ff47f9319
+  depends:
+  - asciitree
+  - fasteners
+  - numcodecs >=0.10.0,!=0.14.0,!=0.14.1
+  - numpy >=1.24
+  - python >=3.11
+  constrains:
+  - ipytree >=0.2.2
+  - notebook
+  - ipywidgets >=8.0.0
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/zarr?source=hash-mapping
+  size: 160435
+  timestamp: 1734108625474
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zeromq-4.3.5-h3b0a872_7.conda
+  sha256: a4dc72c96848f764bb5a5176aa93dd1e9b9e52804137b99daeebba277b31ea10
+  md5: 3947a35e916fcc6b9825449affbf4214
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - krb5 >=1.21.3,<1.22.0a0
+  - libgcc >=13
+  - libsodium >=1.0.20,<1.0.21.0a0
+  - libstdcxx >=13
+  arch: x86_64
+  platform: linux
+  license: MPL-2.0
+  license_family: MOZILLA
+  purls: []
+  size: 335400
+  timestamp: 1731585026517
+- conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.21.0-pyhd8ed1ab_1.conda
+  sha256: 567c04f124525c97a096b65769834b7acb047db24b15a56888a322bf3966c3e1
+  md5: 0c3cc595284c5e8f0f9900a9b228a332
+  depends:
+  - python >=3.9
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/zipp?source=hash-mapping
+  size: 21809
+  timestamp: 1732827613585
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb8e6e7a_0.conda
+  sha256: b8f7b4c7264e84fcedce3929239f5c55e86ae90948c9fdee666f93a70ca58e66
+  md5: 3fe5420e4da42d8110c28f21de6cc127
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=13
+  - libstdcxx >=13
+  - libzlib >=1.3.1,<2.0a0
+  arch: x86_64
+  platform: linux
+  license: BSD-3-Clause
+  purls: []
+  size: 564076
+  timestamp: 1740024448267
diff --git a/pixi.toml b/pixi.toml
new file mode 100644
index 0000000..a330bb3
--- /dev/null
+++ b/pixi.toml
@@ -0,0 +1,52 @@
+[project]
+authors = ["Adrian Seyboldt <adrian.seyboldt@gmail.com>"]
+channels = ["conda-forge"]
+description = "Add a short description here"
+name = "nuts-py"
+platforms = ["linux-64"]
+version = "0.1.0"
+
+[tasks]
+test = "pytest"
+develop = "maturin develop --release"
+get-posteriordb = "git clone 'https://github.com/stan-dev/posteriordb'"
+
+[tasks.bench]
+depends-on = ["develop", "get-posteriordb"]
+cmd = "python -m samplerlab -m posteriordb-fast --posteriordb posteriordb/posterior_database --save-traces --seed 12345"
+
+[dependencies]
+python = ">=3.12.7,<3.13"
+pymc = ">=5.19.0,<6"
+numba = ">=0.60.0,<0.61"
+pytest = ">=8.3.4,<9"
+maturin = ">=1.7.7,<2"
+pip = ">=24.3.1,<25"
+ipykernel = ">=6.29.5,<7"
+seaborn = ">=0.13.2,<0.14"
+threadpoolctl = ">=3.5.0,<4"
+zarr = ">=2.18.3,<3"
+polars = ">=1.16.0,<2"
+viztracer = ">=1.0.0,<2"
+ipywidgets = ">=8.1.5,<9"
+quarto = ">=1.6.40,<2"
+yaml = ">=0.2.5,<0.3"
+pyyaml = ">=6.0.2,<7"
+nbformat = ">=5.10.4,<6"
+nbclient = ">=0.10.2,<0.11"
+cmdstanpy = ">=1.2.5,<2"
+# The jaxlib cuda build seems to be broken around version 0.4.34
+#jax = ">=0.4.35,<0.5"
+
+[pypi-dependencies]
+bridgestan = ">=2.6.0, <3"
+flowjax = { git = "https://github.com/aseyboldt/flowjax.git", rev = "07e7e32217bcfcaa7d68a304f332c925d26ab76f" }
+samplerlab = { git = "https://github.com/aseyboldt/samplerlab/" }
+posteriordb = ">=0.2.0, <0.3"
+#jax = { version = ">=0.5" }
+watermark = ">=2.5.0, <3"
+equinox = "*"
+jaxlib = { version = "*", extras = ["cuda12"] }
+
+[system-requirements]
+#cuda = "12"
diff --git a/pyproject.toml b/pyproject.toml
index 6ef4dff..bd9292f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,17 +2,12 @@
 requires = ["maturin>=1.1,<2.0"]
 build-backend = "maturin"
 
-[tool.maturin]
-module-name = "nutpie._lib"
-python-source = "python"
-features = ["pyo3/extension-module"]
-
 [project]
 name = "nutpie"
 description = "Sample Stan or PyMC models"
 authors = [{ name = "PyMC Developers", email = "pymc.devs@gmail.com" }]
 readme = "README.md"
-requires-python = ">=3.10,<3.13"
+requires-python = ">=3.10,<3.14"
 license = { text = "MIT" }
 classifiers = [
   "Programming Language :: Rust",
@@ -32,56 +27,40 @@ dynamic = ["version"]
 stan = ["bridgestan >= 2.6.1"]
 pymc = ["pymc >= 5.20.1", "numba >= 0.60.0"]
 pymc-jax = ["pymc >= 5.20.1", "jax >= 0.4.27"]
+nnflow = ["flowjax >= 17.0.2"]
+dev = [
+  "bridgestan >= 2.6.1",
+  "pymc >= 5.20.1",
+  "numba >= 0.60.0",
+  "jax >= 0.4.27",
+  "flowjax >= 17.0.2",
+  "pytest",
+]
 all = [
   "bridgestan >= 2.6.1",
   "pymc >= 5.20.1",
   "numba >= 0.60.0",
   "jax >= 0.4.27",
+  "flowjax >= 17.0.2",
 ]
 
 [tool.ruff]
 line-length = 88
-target-version = "py39"
+target-version = "py310"
 show-fixes = true
 output-format = "full"
 
-[tool.ruff.lint]
-select = [
-  "E",   # pycodestyle errors
-  "W",   # pycodestyle warnings
-  "F",   # Pyflakes
-  "I",   # isort
-  "C4",  # flake8-comprehensions
-  "B",   # flake8-bugbear
-  "UP",  # pyupgrade
-  "RUF", # Ruff-specific rules
-  "TID", # flake8-tidy-imports
-  "BLE", # flake8-blind-except
-  "PTH", # flake8-pathlib
-  "A",   # flake8-builtins
-]
-ignore = [
-  "C408", # unnecessary-collection-call (allow dict(a=1, b=2); clarity over speed!)
-  # The following list is recommended to disable these when using ruff's formatter.
-  # (Not all of the following are actually enabled.)
-  "W191",   # tab-indentation
-  "E111",   # indentation-with-invalid-multiple
-  "E114",   # indentation-with-invalid-multiple-comment
-  "E117",   # over-indented
-  "D206",   # indent-with-spaces
-  "D300",   # triple-single-quotes
-  "Q000",   # bad-quotes-inline-string
-  "Q001",   # bad-quotes-multiline-string
-  "Q002",   # bad-quotes-docstring
-  "Q003",   # avoidable-escaped-quote
-  "COM812", # missing-trailing-comma
-  "COM819", # prohibited-trailing-comma
-  "ISC001", # single-line-implicit-string-concatenation
-  "ISC002", # multi-line-implicit-string-concatenation
-]
-
 [tool.ruff.lint.flake8-tidy-imports]
 ban-relative-imports = "all"
 
 [tool.ruff.lint.isort]
 known-first-party = ["nutpie"]
+
+[tool.pyright]
+venvPath = ".pixi/envs/"
+venv = "default"
+
+[tool.maturin]
+module-name = "nutpie._lib"
+python-source = "python"
+features = ["pyo3/extension-module"]
diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py
index 7e344c0..3225d97 100644
--- a/python/nutpie/compile_pymc.py
+++ b/python/nutpie/compile_pymc.py
@@ -6,7 +6,8 @@
 from functools import wraps
 from importlib.util import find_spec
 from math import prod
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, cast
+import threading
 
 import numpy as np
 import pandas as pd
@@ -31,7 +32,7 @@ def intrinsic(f):
     from pytensor.tensor import TensorVariable, Variable
 
 
-def rv_dict_to_flat_array_wrapper(
+def _rv_dict_to_flat_array_wrapper(
     fn: Callable[[SeedType], dict[str, np.ndarray]],
     names: list[str],
     shapes: list[tuple[int]],
@@ -273,7 +274,7 @@ def _compile_pymc_model_numba(
         warnings.filterwarnings(
             "ignore",
             message="Cannot cache compiled function .* as it uses dynamic globals",
-            category=numba.NumbaWarning,
+            category=numba.NumbaWarning,  # type: ignore
         )
 
         logp_numba = numba.cfunc(c_sig, **kwargs)(logp_numba_raw)
@@ -286,7 +287,7 @@ def _compile_pymc_model_numba(
         warnings.filterwarnings(
             "ignore",
             message="Cannot cache compiled function .* as it uses dynamic globals",
-            category=numba.NumbaWarning,
+            category=numba.NumbaWarning,  # type: ignore
         )
 
         expand_numba = numba.cfunc(c_sig_expand, **kwargs)(expand_numba_raw)
@@ -376,14 +377,24 @@ def _compile_pymc_model_jax(
     logp_fn = logp_fn_pt.vm.jit_fn
     expand_fn = expand_fn_pt.vm.jit_fn
 
+    logp_shared_names = [var.name for var in logp_fn_pt.get_shared()]
+    expand_shared_names = [var.name for var in expand_fn_pt.get_shared()]
+
     if gradient_backend == "jax":
         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)
 
+        # static_argnums = list(range(1, len(logp_shared_names) + 1))
+        logp_fn_jax_grad = jax.jit(
+            logp_fn_jax_grad,
+            # static_argnums=static_argnums,
+        )
+
         logp_fn = logp_fn_jax_grad
+    else:
+        orig_logp_fn = None
 
     shared_data = {}
     shared_vars = {}
@@ -395,9 +406,6 @@ def logp_fn_jax_grad(x, *shared):
         shared_vars[val.name] = val
         seen.add(val)
 
-    logp_shared_names = [var.name for var in logp_fn_pt.get_shared()]
-    expand_shared_names = [var.name for var in expand_fn_pt.get_shared()]
-
     def make_logp_func():
         def logp(x, **shared):
             logp, grad = logp_fn(x, *[shared[name] for name in logp_shared_names])
@@ -406,7 +414,8 @@ def logp(x, **shared):
         return logp
 
     names, slices, shapes = shape_info
-    dtypes = [np.float64] * len(names)
+    # TODO do not cast to float64
+    dtypes = [np.dtype("float64")] * len(names)
 
     def make_expand_func(seed1, seed2, chain):
         # TODO handle seeds
@@ -432,6 +441,7 @@ def expand(x, **shared):
         shared_data=shared_data,
         dims=dims,
         coords=coords,
+        raw_logp_fn=orig_logp_fn,
     )
 
 
@@ -509,6 +519,8 @@ def compile_pymc_model(
         return_transformed=True,
     )
 
+    initial_point_fn = _wrap_with_lock(initial_point_fn)
+
     if backend.lower() == "numba":
         if gradient_backend == "jax":
             raise ValueError("Gradient backend cannot be jax when using numba backend")
@@ -530,7 +542,18 @@ def compile_pymc_model(
         raise ValueError(f"Backend must be one of numba and jax. Got {backend}")
 
 
-def _compute_shapes(model):
+def _wrap_with_lock(func: Callable) -> Callable:
+    lock = threading.Lock()
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with lock:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def _compute_shapes(model) -> dict[str, tuple[int, ...]]:
     import pytensor
     from pymc.initial_point import make_initial_point_fn
 
@@ -663,7 +686,7 @@ def _make_functions(
 
     num_free_vars = count
 
-    initial_point_fn = rv_dict_to_flat_array_wrapper(
+    initial_point_fn = _rv_dict_to_flat_array_wrapper(
         pymc_initial_point_fn, names=joined_names, shapes=joined_shapes
     )
 
@@ -712,7 +735,7 @@ def _make_functions(
 
     for var in remaining_rvs:
         all_names.append(var.name)
-        shape = shapes[var.name]
+        shape = cast(tuple[int, ...], shapes[var.name])
         all_shapes.append(shape)
         length = prod(shape)
         all_slices.append(slice(count, count + length))
diff --git a/python/nutpie/compile_stan.py b/python/nutpie/compile_stan.py
index 7a28052..138652d 100644
--- a/python/nutpie/compile_stan.py
+++ b/python/nutpie/compile_stan.py
@@ -1,6 +1,7 @@
 import json
 import tempfile
 from dataclasses import dataclass, replace
+from functools import partial
 from importlib.util import find_spec
 from pathlib import Path
 from typing import Any, Optional
@@ -11,6 +12,7 @@
 
 from nutpie import _lib
 from nutpie.sample import CompiledModel
+from nutpie.transform_adapter import make_transform_adapter
 
 
 class _NumpyArrayEncoder(json.JSONEncoder):
@@ -28,6 +30,7 @@ class CompiledStanModel(CompiledModel):
     library: Any
     model: Any
     model_name: Optional[str] = None
+    _transform_adapt_args: dict | None = None
 
     def with_data(self, *, seed=None, **updates):
         if self.data is None:
@@ -42,7 +45,15 @@ def with_data(self, *, seed=None, **updates):
         else:
             data_json = None
 
-        model = _lib.StanModel(self.library, seed, data_json)
+        kwargs = self._transform_adapt_args
+        if kwargs is None:
+            kwargs = {}
+        make_adapter = partial(
+            make_transform_adapter(**kwargs),
+            logp_fn=None,
+        )
+
+        model = _lib.StanModel(self.library, seed, data_json, make_adapter)
         coords = self._coords
         if coords is None:
             coords = {}
@@ -75,6 +86,9 @@ def with_dims(self, **dims):
         dims_new.update(dims)
         return replace(self, dims=dims_new)
 
+    def with_transform_adapt(self, **kwargs):
+        return replace(self, _transform_adapt_args=kwargs).with_data()
+
     def _make_model(self, init_mean):
         if self.model is None:
             return self.with_data().model
diff --git a/python/nutpie/compiled_pyfunc.py b/python/nutpie/compiled_pyfunc.py
index 9ede109..d0f9732 100644
--- a/python/nutpie/compiled_pyfunc.py
+++ b/python/nutpie/compiled_pyfunc.py
@@ -5,8 +5,9 @@
 
 import numpy as np
 
-from nutpie import _lib
+from nutpie import _lib  # type: ignore
 from nutpie.sample import CompiledModel
+from nutpie.transform_adapter import make_transform_adapter
 
 SeedType = int
 
@@ -20,6 +21,8 @@ class PyFuncModel(CompiledModel):
     _n_dim: int
     _variables: list[_lib.PyVariable]
     _coords: dict[str, Any]
+    _raw_logp_fn: Callable | None
+    _transform_adapt_args: dict | None = None
 
     @property
     def shapes(self) -> dict[str, tuple[int, ...]]:
@@ -42,6 +45,9 @@ def with_data(self, **updates):
         updated.update(**updates)
         return dataclasses.replace(self, _shared_data=updated)
 
+    def with_transform_adapt(self, **kwargs):
+        return dataclasses.replace(self, _transform_adapt_args=kwargs)
+
     def _make_sampler(self, settings, init_mean, cores, progress_type):
         model = self._make_model(init_mean)
         return _lib.PySampler.from_pyfunc(
@@ -60,12 +66,24 @@ def make_expand_func(seed1, seed2, chain):
             expand_fn = self._make_expand_func(seed1, seed2, chain)
             return partial(expand_fn, **self._shared_data)
 
+        if self._raw_logp_fn is not None:
+            kwargs = self._transform_adapt_args
+            if kwargs is None:
+                kwargs = {}
+            make_adapter = partial(
+                make_transform_adapter(**kwargs),
+                logp_fn=self._raw_logp_fn,
+            )
+        else:
+            make_adapter = None
+
         return _lib.PyModel(
             make_logp_func,
             make_expand_func,
             self._variables,
             self.n_dim,
-            self._make_initial_points,
+            init_point_func=self._make_initial_points,
+            transform_adapter=make_adapter,
         )
 
 
@@ -81,6 +99,8 @@ def from_pyfunc(
     dims: dict[str, tuple[str, ...]] | None = None,
     shared_data: dict[str, Any] | None = None,
     make_initial_point_fn: Callable[[SeedType], np.ndarray] | None,
+    make_transform_adapter=None,
+    raw_logp_fn=None,
 ):
     variables = []
     for name, shape, dtype in zip(
@@ -111,4 +131,5 @@ def from_pyfunc(
         _make_initial_points=make_initial_point_fn,
         _variables=variables,
         _shared_data=shared_data,
+        _raw_logp_fn=raw_logp_fn,
     )
diff --git a/python/nutpie/normalizing_flow.py b/python/nutpie/normalizing_flow.py
new file mode 100644
index 0000000..631af8b
--- /dev/null
+++ b/python/nutpie/normalizing_flow.py
@@ -0,0 +1,919 @@
+from typing import Union, Literal, Callable
+import math
+import itertools
+
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from flowjax import bijections
+import flowjax.distributions
+import flowjax.flows
+import numpy as np
+from paramax import Parameterize
+
+
+def _generate_sequences(k, r_vals):
+    """
+    Generate all binary sequences of length k with exactly r 1's.
+    The sequences are stored in a preallocated boolean NumPy array of shape (N, k),
+    where N = comb(k, r). A True value represents a '1' and False represents a '0'.
+
+    Parameters:
+        k (int): The length of each sequence.
+        r (int): The exact number of ones in each sequence.
+
+    Returns:
+        A NumPy boolean array of shape (comb(k, r), k) containing all sequences.
+    """
+    all_sequences = []
+    for r in r_vals:
+        N = math.comb(k, r)  # number of sequences
+        sequences = np.zeros((N, k), dtype=bool)
+        # Use enumerate on all combinations where ones appear.
+        for i, ones_positions in enumerate(itertools.combinations(range(k), r)):
+            sequences[i, list(ones_positions)] = True
+        all_sequences.append(sequences)
+    return np.concatenate(all_sequences, axis=0)
+
+
+def _max_run_length(seq):
+    """
+    Given a 1D boolean NumPy array 'seq', compute the maximum run length of consecutive
+    identical values (either True or False).
+
+    Parameters:
+        seq (np.array): A 1D boolean array.
+
+    Returns:
+        The length (int) of the longest run.
+    """
+    # If the sequence is empty, return 0.
+    if seq.size == 0:
+        return 0
+
+    # Convert boolean to int (0 or 1) so we can use np.diff.
+    arr = seq.astype(int)
+    # Compute differences between consecutive elements.
+    diffs = np.diff(arr)
+    # Positions where the value changes:
+    change_indices = np.nonzero(diffs)[0]
+
+    if change_indices.size == 0:
+        # No changes at all, so the entire sequence is one run.
+        return seq.size
+
+    # To compute the run lengths, add the "start" index (-1) and the last index.
+    # For example, if change_indices = [i1, i2, ..., in],
+    # then the runs are: (i1 - (-1)), (i2 - i1), ..., (seq.size-1 - in).
+    boundaries = np.concatenate(([-1], change_indices, [seq.size - 1]))
+    run_lengths = np.diff(boundaries)
+    return int(run_lengths.max())
+
+
+def _filter_sequences(sequences, m):
+    """
+    Filter a 2D NumPy boolean array 'sequences' (each row a binary sequence) so that
+    only sequences with maximum run length (of 0's or 1's) at most m are kept.
+
+    Parameters:
+        sequences (np.array): A 2D boolean array of shape (N, k).
+        m (int): Maximum allowed run length.
+
+    Returns:
+        A NumPy array containing only the rows (sequences) that pass the filter.
+    """
+    filtered = []
+    for seq in sequences:
+        if _max_run_length(seq) <= m:
+            filtered.append(seq)
+    return np.array(filtered)
+
+
+def _generate_permutations(rng, n_dim, n_layers, max_run=3):
+    if n_layers == 1:
+        r = [0, 1]
+    elif n_layers == 2:
+        r = [1]
+    else:
+        if n_layers % 2 == 0:
+            half = n_layers // 2
+            r = [half - 1, half, half + 1]
+        else:
+            half = n_layers // 2
+            r = [half, half + 1]
+
+    all_sequences = _generate_sequences(n_layers, r)
+    valid_sequences = _filter_sequences(all_sequences, max_run)
+
+    valid_sequences = np.repeat(
+        valid_sequences, n_dim // len(valid_sequences) + 1, axis=0
+    )
+    rng.shuffle(valid_sequences, axis=0)
+    is_in_first = valid_sequences[:n_dim]
+    rng = np.random.default_rng(42)
+    permutations = (~is_in_first).argsort(axis=0, kind="stable")
+    return permutations.T, is_in_first.sum(0)
+
+
+# Fix upstream bug for zero-dimensional arrays
+class Linear(eqx.Module, strict=True):
+    """Performs a linear transformation."""
+
+    weight: jax.Array
+    bias: jax.Array | None
+    in_features: Union[int, Literal["scalar"]] = eqx.field(static=True)
+    out_features: Union[int, Literal["scalar"]] = eqx.field(static=True)
+    use_bias: bool = eqx.field(static=True)
+
+    def __init__(
+        self,
+        in_features: Union[int, Literal["scalar"]],
+        out_features: Union[int, Literal["scalar"]],
+        use_bias: bool = True,
+        dtype=None,
+        *,
+        key,
+    ):
+        """**Arguments:**
+
+        - `in_features`: The input size. The input to the layer should be a vector of
+            shape `(in_features,)`
+        - `out_features`: The output size. The output from the layer will be a vector
+            of shape `(out_features,)`.
+        - `use_bias`: Whether to add on a bias as well.
+        - `dtype`: The dtype to use for the weight and the bias in this layer.
+            Defaults to either `jax.numpy.float32` or `jax.numpy.float64` depending
+            on whether JAX is in 64-bit mode.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+
+        Note that `in_features` also supports the string `"scalar"` as a special value.
+        In this case the input to the layer should be of shape `()`.
+
+        Likewise `out_features` can also be a string `"scalar"`, in which case the
+        output from the layer will have shape `()`.
+        """
+        dtype = np.float32 if dtype is None else dtype
+        wkey, bkey = jax.random.split(key, 2)
+        in_features_ = 1 if in_features == "scalar" else in_features
+        out_features_ = 1 if out_features == "scalar" else out_features
+        if in_features_ == 0:
+            lim = 1.0
+        else:
+            lim = 1 / math.sqrt(in_features_)
+        wshape = (out_features_, in_features_)
+        self.weight = eqx.nn._misc.default_init(wkey, wshape, dtype, lim)
+        bshape = (out_features_,)
+        self.bias = (
+            eqx.nn._misc.default_init(bkey, bshape, dtype, lim) if use_bias else None
+        )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = use_bias
+
+    @jax.named_scope("eqx.nn.Linear")
+    def __call__(self, x: jax.Array, *, key=None) -> jax.Array:
+        """**Arguments:**
+
+        - `x`: The input. Should be a JAX array of shape `(in_features,)`. (Or shape
+            `()` if `in_features="scalar"`.)
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        !!! info
+
+            If you want to use higher order tensors as inputs (for example featuring "
+            "batch dimensions) then use `jax.vmap`. For example, for an input `x` of "
+            "shape `(batch, in_features)`, using
+            ```python
+            linear = equinox.nn.Linear(...)
+            jax.vmap(linear)(x)
+            ```
+            will produce the appropriate output of shape `(batch, out_features)`.
+
+        **Returns:**
+
+        A JAX array of shape `(out_features,)`. (Or shape `()` if
+        `out_features="scalar"`.)
+        """
+
+        if self.in_features == "scalar":
+            if jnp.shape(x) != ():
+                raise ValueError("x must have scalar shape")
+            x = jnp.broadcast_to(x, (1,))
+        x = self.weight @ x
+        if self.bias is not None:
+            x = x + self.bias
+        if self.out_features == "scalar":
+            assert jnp.shape(x) == (1,)
+            x = jnp.squeeze(x)
+        return x
+
+
+class FactoredMLP(eqx.Module, strict=True):
+    """Standard Multi-Layer Perceptron; also known as a feed-forward network.
+
+    !!! faq
+
+        If you get a TypeError saying an object is not a valid JAX type, see the
+            [FAQ](https://docs.kidger.site/equinox/faq/)."""
+
+    layers: tuple[tuple[Linear, Linear], ...]
+    activation: tuple[Callable, ...]
+    final_activation: Callable
+    use_bias: bool = eqx.field(static=True)
+    use_final_bias: bool = eqx.field(static=True)
+    in_size: Union[int, Literal["scalar"]] = eqx.field(static=True)
+    out_size: Union[int, Literal["scalar"]] = eqx.field(static=True)
+    width_size: tuple[int, ...] = eqx.field(static=True)
+    depth: int = eqx.field(static=True)
+
+    def __init__(
+        self,
+        in_size: Union[int, Literal["scalar"]],
+        out_size: Union[int, Literal["scalar"]],
+        width_size: int | tuple[int | tuple[int, int], ...],
+        depth: int,
+        activation: Callable = jax.nn.relu,
+        final_activation: Callable = lambda x: x,
+        use_bias: bool = True,
+        use_final_bias: bool = True,
+        dtype=None,
+        *,
+        key,
+    ):
+        """**Arguments**:
+
+        - `in_size`: The input size. The input to the module should be a vector of
+            shape `(in_features,)`
+        - `out_size`: The output size. The output from the module will be a vector
+            of shape `(out_features,)`.
+        - `width_size`: The size of each hidden layer.
+        - `depth`: The number of hidden layers, including the output layer.
+            For example, `depth=2` results in an network with layers:
+            [`Linear(in_size, width_size)`, `Linear(width_size, width_size)`,
+            `Linear(width_size, out_size)`].
+        - `activation`: The activation function after each hidden layer. Defaults to
+            ReLU.
+        - `final_activation`: The activation function after the output layer. Defaults
+            to the identity.
+        - `use_bias`: Whether to add on a bias to internal layers. Defaults
+            to `True`.
+        - `use_final_bias`: Whether to add on a bias to the final layer. Defaults
+            to `True`.
+        - `dtype`: The dtype to use for all the weights and biases in this MLP.
+            Defaults to either `jax.numpy.float32` or `jax.numpy.float64` depending
+            on whether JAX is in 64-bit mode.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+
+        Note that `in_size` also supports the string `"scalar"` as a special value.
+        In this case the input to the module should be of shape `()`.
+
+        Likewise `out_size` can also be a string `"scalar"`, in which case the
+        output from the module will have shape `()`.
+        """
+        keys = jax.random.split(key, depth + 1)
+        layers = []
+        if isinstance(width_size, int):
+            width_size = (width_size,) * depth
+
+        assert len(width_size) == depth
+        activations: list[Callable] = []
+
+        if depth == 0:
+            layers.append(
+                Linear(in_size, out_size, use_final_bias, dtype=dtype, key=keys[0])
+            )
+        else:
+            if isinstance(width_size[0], tuple):
+                n, k = width_size[0]
+                key1, key2 = jax.random.split(keys[0])
+                U = Linear(in_size, n, use_bias=False, dtype=dtype, key=key1)
+                K = Linear(n, k, use_bias=True, dtype=dtype, key=key2)
+                layers.append((U, K))
+            else:
+                k = width_size[0]
+                layers.append(Linear(in_size, k, use_bias, dtype=dtype, key=keys[0]))
+            activations.append(eqx.filter_vmap(lambda: activation, axis_size=k)())
+
+            for i in range(depth - 1):
+                if isinstance(width_size[i + 1], tuple):
+                    n, k_new = width_size[i + 1]
+                    key1, key2 = jax.random.split(keys[i + 1])
+                    U = Linear(k, n, use_bias=False, dtype=dtype, key=key1)
+                    K = Linear(n, k_new, use_bias=True, dtype=dtype, key=key2)
+                    layers.append((U, K))
+                    k = k_new
+                else:
+                    layers.append(
+                        Linear(
+                            k, width_size[i + 1], use_bias, dtype=dtype, key=keys[i + 1]
+                        )
+                    )
+                    k = width_size[i + 1]
+                activations.append(eqx.filter_vmap(lambda: activation, axis_size=k)())
+
+            if isinstance(out_size, tuple):
+                n, k_new = out_size
+                key1, key2 = jax.random.split(keys[-1])
+                U = Linear(k, n, use_bias=False, dtype=dtype, key=key1)
+                K = Linear(n, k_new, use_bias=True, dtype=dtype, key=key2)
+                k = k_new
+                layers.append((U, K))
+            else:
+                layers.append(
+                    Linear(k, out_size, use_final_bias, dtype=dtype, key=keys[-1])
+                )
+        self.layers = tuple(layers)
+        self.in_size = in_size
+        self.out_size = out_size
+        self.width_size = width_size
+        self.depth = depth
+        # In case `activation` or `final_activation` are learnt, then make a separate
+        # copy of their weights for every neuron.
+        self.activation = tuple(activations)
+        if out_size == "scalar":
+            self.final_activation = final_activation
+        else:
+            self.final_activation = eqx.filter_vmap(
+                lambda: final_activation, axis_size=out_size
+            )()
+        self.use_bias = use_bias
+        self.use_final_bias = use_final_bias
+
+    @jax.named_scope("eqx.nn.MLP")
+    def __call__(self, x: jax.Array, *, key=None) -> jax.Array:
+        """**Arguments:**
+
+        - `x`: A JAX array with shape `(in_size,)`. (Or shape `()` if
+            `in_size="scalar"`.)
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        A JAX array with shape `(out_size,)`. (Or shape `()` if `out_size="scalar"`.)
+        """
+        for i, (layer, act) in enumerate(zip(self.layers[:-1], self.activation)):
+            if isinstance(layer, tuple):
+                U, K = layer
+                x = U(x)
+                x = K(x)
+            else:
+                x = layer(x)
+            layer_activation = jax.tree.map(
+                lambda x: x[i] if eqx.is_array(x) else x, act
+            )
+            x = eqx.filter_vmap(lambda a, b: a(b))(layer_activation, x)
+
+        if isinstance(self.layers[-1], tuple):
+            U, K = self.layers[-1]
+            x = U(x)
+            x = K(x)
+        else:
+            x = self.layers[-1](x)
+
+        if self.out_size == "scalar":
+            x = self.final_activation(x)
+        else:
+            x = eqx.filter_vmap(lambda a, b: a(b))(self.final_activation, x)
+        return x
+
+
+def make_mvscale(key, n_dim, size, randomize_base=False):
+    def make_single_hh(key, idx):
+        key1, key2 = jax.random.split(key)
+        params = jax.random.normal(key1, (n_dim,))
+        params = params / jnp.linalg.norm(params)
+        mvscale = bijections.MvScale(params)
+        return mvscale
+
+    keys = jax.random.split(key, size)
+
+    if randomize_base:
+        key, key_base = jax.random.split(key)
+        indices = jax.random.randint(key_base, (size,), 0, n_dim)
+    else:
+        indices = [val % n_dim for val in range(size)]
+
+    return bijections.Chain(
+        [make_single_hh(key, idx) for key, idx in zip(keys, indices)]
+    )
+
+
+def make_hh(key, n_dim, size, randomize_base=False):
+    def make_single_hh(key, idx):
+        key1, key2 = jax.random.split(key)
+        params = jax.random.normal(key1, (n_dim,)) * 1e-2
+        return bijections.Householder(params, base_index=idx)
+
+    keys = jax.random.split(key, size)
+
+    if randomize_base:
+        key, key_base = jax.random.split(key)
+        indices = jax.random.randint(key_base, (size,), 0, n_dim)
+    else:
+        indices = [val % n_dim for val in range(size)]
+
+    return bijections.Chain(
+        [make_single_hh(key, idx) for key, idx in zip(keys, indices)]
+    )
+
+
+def make_elemwise_trafo(key, n_dim, *, count=1):
+    def make_elemwise(key, loc):
+        key1, key2 = jax.random.split(key)
+        scale = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(()))
+        theta = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(()))
+
+        affine = bijections.AsymmetricAffine(
+            loc,
+            jnp.ones(()),
+            jnp.ones(()),
+        )
+
+        affine = eqx.tree_at(
+            where=lambda aff: aff.scale,
+            pytree=affine,
+            replace=scale,
+        )
+        affine = eqx.tree_at(
+            where=lambda aff: aff.theta,
+            pytree=affine,
+            replace=theta,
+        )
+
+        return affine
+
+    def make(key):
+        keys = jax.random.split(key, count + 1)
+        key, keys = keys[0], keys[1:]
+        loc = jax.random.normal(key=key, shape=(count,)) * 2
+        loc = loc - loc.mean()
+        return bijections.Chain([make_elemwise(key, mu) for key, mu in zip(keys, loc)])
+
+    keys = jax.random.split(key, n_dim)
+    make_affine = eqx.filter_vmap(make, axis_size=n_dim)(keys)
+    return bijections.Vmap(make_affine, in_axes=eqx.if_array(0))
+
+
+def make_elemwise_trafo_(key, n_dim, *, count=1):
+    def make_elemwise(key):
+        scale = Parameterize(
+            lambda x: x + jnp.sqrt(1 + x**2),
+            jax.random.normal(key=key) / 5,
+        )
+        theta = Parameterize(
+            lambda x: x + jnp.sqrt(1 + x**2),
+            jax.random.normal(key=key) / 5,
+        )
+
+        affine = bijections.AsymmetricAffine(
+            jax.random.normal(key=key) * 2,
+            jnp.ones(()),
+            jnp.ones(()),
+        )
+
+        affine = eqx.tree_at(
+            where=lambda aff: aff.scale,
+            pytree=affine,
+            replace=scale,
+        )
+        affine = eqx.tree_at(
+            where=lambda aff: aff.theta,
+            pytree=affine,
+            replace=theta,
+        )
+
+        return affine
+
+    def make(key):
+        keys = jax.random.split(key, count)
+        return bijections.Scan(eqx.filter_vmap(make_elemwise)(keys))
+
+    keys = jax.random.split(key, n_dim)
+    make_affine = eqx.filter_vmap(make)(keys)
+    return bijections.Vmap(make_affine())
+
+
+def make_coupling(key, dim, n_untransformed, **kwargs):
+    n_transformed = dim - n_untransformed
+
+    mvscale = make_mvscale(key, n_transformed, 1, randomize_base=True)
+
+    nn_width = kwargs.get("nn_width", None)
+
+    if nn_width is None:
+        if dim > 128:
+            nn_width = (64, 2 * dim)
+        else:
+            nn_width = 2 * dim
+
+    transformer = bijections.Chain(
+        [
+            make_elemwise_trafo(key, n_transformed, count=3),
+            #mvscale,
+        ]
+    )
+
+    def make_mlp(out_size):
+        if isinstance(nn_width, int):
+            width = [nn_width]
+        else:
+            width = nn_width
+        if isinstance(width[0], tuple):
+            out = (width[0][0], out_size)
+        else:
+            out = out_size
+
+        return FactoredMLP(
+            n_untransformed,
+            out,
+            width,
+            depth=1,
+            key=key,
+            dtype=jnp.float32,
+            activation=jax.nn.gelu,
+        )
+
+    return bijections.Coupling(
+        key,
+        transformer=transformer,
+        untransformed_dim=n_untransformed,
+        dim=dim,
+        conditioner=make_mlp,
+        nn_depth=1,
+        **kwargs,
+    )
+
+
+def make_flow(
+    seed,
+    positions,
+    gradients,
+    *,
+    zero_init=False,
+    householder_layer=False,
+    dct_layer=False,
+    untransformed_dim: int | list[int | None] | None = None,
+    n_layers,
+    nn_width=None,
+):
+    from flowjax import bijections
+
+    positions = np.array(positions)
+    gradients = np.array(gradients)
+
+    if len(positions) == 0:
+        return
+
+    n_draws, n_dim = positions.shape
+
+    if n_dim < 2:
+        n_layers = 0
+
+    assert positions.shape == gradients.shape
+
+    if n_draws == 0:
+        raise ValueError("No draws")
+    elif n_draws == 1:
+        assert np.all(gradients != 0)
+        diag = np.clip(1 / jnp.sqrt(jnp.abs(gradients[0])), 1e-5, 1e5)
+        assert np.isfinite(diag).all()
+        mean = jnp.zeros_like(diag)
+    else:
+        pos_std = np.clip(positions.std(0), 1e-8, 1e8)
+        grad_std = np.clip(gradients.std(0), 1e-8, 1e8)
+        diag = jnp.sqrt(pos_std / grad_std)
+        mean = positions.mean(0) + gradients.mean(0) * diag * diag
+
+    key = jax.random.PRNGKey(seed % (2**63))
+
+    diag_param = Parameterize(
+        lambda x: x + jnp.sqrt(1 + x**2),
+        (diag**2 - 1) / (2 * diag),
+    )
+    diag_affine = bijections.Affine(mean, diag)
+    diag_affine = eqx.tree_at(
+        where=lambda aff: aff.scale,
+        pytree=diag_affine,
+        replace=diag_param,
+    )
+
+    flows = [
+        diag_affine,
+    ]
+
+    if n_layers == 0:
+        return bijections.Chain(flows)
+
+    def make_layer(key, untransformed_dim: int | None, permutation=None):
+        key, key_couple, key_permute, key_hh = jax.random.split(key, 4)
+
+        if untransformed_dim is None:
+            untransformed_dim = n_dim // 2
+
+        if untransformed_dim < 0:
+            untransformed_dim = n_dim + untransformed_dim
+
+        coupling = make_coupling(
+            key_couple,
+            n_dim,
+            untransformed_dim,
+            nn_activation=jax.nn.gelu,
+            nn_width=nn_width,
+        )
+
+        if zero_init:
+            coupling = jax.tree_util.tree_map(
+                lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
+                coupling,
+            )
+
+        flow = coupling
+
+        if householder_layer:
+            hh = make_hh(key_hh, n_dim, 1, randomize_base=False)
+            flow = bijections.Sandwich(hh, flow)
+
+        def add_default_permute(bijection, dim, key):
+            if dim == 1:
+                return bijection
+            if dim == 2:
+                outer = bijections.Flip((dim,))
+            else:
+                outer = bijections.Permute(jax.random.permutation(key, jnp.arange(dim)))
+
+            return bijections.Sandwich(outer, bijection)
+
+        if permutation is None:
+            flow = add_default_permute(flow, n_dim, key_permute)
+        else:
+            flow = bijections.Sandwich(bijections.Permute(permutation), flow)
+
+        mvscale = make_mvscale(key, n_dim, 1, randomize_base=True)
+
+        flow = bijections.Chain(
+            [
+                mvscale,
+                flow,
+            ]
+        )
+
+        return flow
+
+    key, key_permute = jax.random.split(key)
+    keys = jax.random.split(key, n_layers)
+
+    if untransformed_dim is None:
+        # TODO better rng?
+        rng = np.random.default_rng(int(jax.random.randint(key, (), 0, 2**30)))
+        permutation, lengths = _generate_permutations(rng, n_dim, n_layers)
+        layers = []
+        for i, (key, p, length) in enumerate(zip(keys, permutation, lengths)):
+            layers.append(make_layer(key, int(length), p))
+        bijection = bijections.Chain(layers)
+    elif isinstance(untransformed_dim, int):
+        make_layers = eqx.filter_vmap(make_layer)
+        layers = make_layers(keys, untransformed_dim)
+        bijection = bijections.Scan(layers)
+    else:
+        layers = []
+        for i, (key, num_untrafo) in enumerate(zip(keys, untransformed_dim)):
+            if i % 2 == 0 or not dct_layer:
+                layers.append(make_layer(key, num_untrafo))
+            else:
+                inner = make_layer(key, num_untrafo)
+                outer = bijections.DCT(inner.shape)
+
+                layers.append(bijections.Sandwich(outer, inner))
+
+        bijection = bijections.Chain(layers)
+
+    return bijections.Chain([bijection, *flows])
+
+
+def extend_flow(
+    key,
+    base,
+    loss_fn,
+    positions,
+    gradients,
+    logps,
+    layer: int,
+    *,
+    extension_var_count=4,
+    zero_init=False,
+    householder_layer=False,
+    untransformed_dim: int | list[int | None] | None = None,
+    dct: bool = False,
+    extension_var_trafo_count=2,
+    verbose: bool = False,
+    nn_width=None,
+    nn_depth=None,
+):
+    n_draws, n_dim = positions.shape
+
+    if n_dim < 2:
+        return base
+
+    if n_dim <= extension_var_count:
+        extension_var_count = n_dim - 1
+        extension_var_trafo_count = 1
+
+    if dct:
+        flow = flowjax.flows.Transformed(
+            flowjax.distributions.StandardNormal(base.shape),
+            bijections.Chain([bijections.DCT(shape=(n_dim,)), base]),
+        )
+    else:
+        flow = flowjax.flows.Transformed(
+            flowjax.distributions.StandardNormal(base.shape), base
+        )
+
+    params, static = eqx.partition(flow, eqx.is_inexact_array)
+    costs = loss_fn(
+        params,
+        static,
+        positions,
+        gradients,
+        logps,
+        return_elemwise_costs=True,
+    )
+
+    if verbose:
+        print(max(costs), costs)
+        print("dct:", dct)
+    idxs = np.argsort(costs)
+
+    permute = bijections.Permute(idxs)
+
+    if True:
+        scale = Parameterize(
+            lambda x: x + jnp.sqrt(1 + x**2),
+            jnp.array(0.0),
+        )
+        theta = Parameterize(
+            lambda x: x + jnp.sqrt(1 + x**2),
+            jnp.array(0.0),
+        )
+
+        affine = bijections.AsymmetricAffine(jnp.zeros(()), jnp.ones(()), jnp.ones(()))
+
+        affine = eqx.tree_at(
+            where=lambda aff: aff.scale,
+            pytree=affine,
+            replace=scale,
+        )
+        affine = eqx.tree_at(
+            where=lambda aff: aff.theta,
+            pytree=affine,
+            replace=theta,
+        )
+
+        do_flip = layer % 2 == 0
+
+        if nn_width is None:
+            width = 16
+        else:
+            width = nn_width
+
+        if do_flip:
+            coupling = bijections.coupling.Coupling(
+                key,
+                transformer=affine,
+                untransformed_dim=n_dim - extension_var_trafo_count,
+                dim=n_dim,
+                nn_activation=jax.nn.gelu,
+                nn_width=width,
+                nn_depth=nn_depth,
+            )
+
+            inner_permute = bijections.Permute(
+                jnp.concatenate(
+                    [
+                        jnp.arange(n_dim - extension_var_count),
+                        jax.random.permutation(
+                            key, jnp.arange(n_dim - extension_var_count, n_dim)
+                        ),
+                    ]
+                )
+            )
+        else:
+            coupling = bijections.coupling.Coupling(
+                key,
+                transformer=affine,
+                untransformed_dim=extension_var_trafo_count,
+                dim=n_dim,
+                nn_activation=jax.nn.gelu,
+                nn_width=width,
+                nn_depth=nn_depth,
+            )
+
+            inner_permute = bijections.Permute(
+                jnp.concatenate(
+                    [
+                        jax.random.permutation(
+                            key, jnp.arange(n_dim - extension_var_count, n_dim)
+                        ),
+                        jnp.arange(n_dim - extension_var_count),
+                    ]
+                )
+            )
+
+        if zero_init:
+            coupling = jax.tree_util.tree_map(
+                lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
+                coupling,
+            )
+
+        inner = bijections.Sandwich(inner_permute, coupling)
+
+        if False:
+            scale = Parameterize(
+                lambda x: x + jnp.sqrt(1 + x**2),
+                jnp.array(0.0),
+            )
+            affine = eqx.tree_at(
+                where=lambda aff: aff.scale,
+                pytree=flowjax.bijections.Affine(),
+                replace=scale,
+            )
+
+            if nn_width is None:
+                width = 16
+            else:
+                width = nn_width
+
+            coupling = flowjax.bijections.coupling.Coupling(
+                key,
+                transformer=affine,
+                untransformed_dim=extension_var_trafo_count,
+                dim=n_dim,
+                nn_activation=jax.nn.gelu,
+                nn_width=width,
+                nn_depth=nn_depth,
+            )
+
+            if zero_init:
+                coupling = jax.tree_util.tree_map(
+                    lambda x: x * 1e-3 if eqx.is_inexact_array(x) else x,
+                    coupling,
+                )
+
+            if verbose:
+                print(costs[permute.permutation][inner.outer.permutation])
+
+            inner = bijections.Sandwich(
+                inner.outer,
+                bijections.Chain(
+                    [
+                        bijections.Sandwich(bijections.Flip(shape=(n_dim,)), coupling),
+                        inner.inner,
+                    ]
+                ),
+            )
+
+    if dct:
+        new_layer = bijections.Sandwich(
+            bijections.DCT(shape=(n_dim,)),
+            bijections.Sandwich(permute, inner),
+        )
+    else:
+        new_layer = bijections.Sandwich(permute, inner)
+
+    scale = Parameterize(
+        lambda x: x + jnp.sqrt(1 + x**2),
+        jnp.zeros(n_dim),
+    )
+    affine = eqx.tree_at(
+        where=lambda aff: aff.scale,
+        pytree=bijections.Affine(jnp.zeros(n_dim), jnp.ones(n_dim)),
+        replace=scale,
+    )
+
+    pre = []
+    if layer % 2 == 0:
+        pre.append(bijections.Neg(shape=(n_dim,)))
+
+    nonlin_layer = bijections.Sandwich(
+        bijections.Chain(
+            [
+                *pre,
+                bijections.Vmap(bijections.SoftPlusX(), axis_size=n_dim),
+            ]
+        ),
+        affine,
+    )
+    scale = Parameterize(
+        lambda x: x + jnp.sqrt(1 + x**2),
+        jnp.zeros(n_dim),
+    )
+    affine = eqx.tree_at(
+        where=lambda aff: aff.scale,
+        pytree=bijections.Affine(jnp.zeros(n_dim), jnp.ones(n_dim)),
+        replace=scale,
+    )
+    return bijections.Chain([new_layer, nonlin_layer, affine, base])
diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py
index 6c5bfc6..356b8d0 100644
--- a/python/nutpie/sample.py
+++ b/python/nutpie/sample.py
@@ -1,13 +1,13 @@
 import os
 from dataclasses import dataclass
-from typing import Any, Literal, Optional, overload
+from typing import Any, Literal, Optional, cast, overload
 
 import arviz
 import numpy as np
 import pandas as pd
 import pyarrow
 
-from nutpie import _lib
+from nutpie import _lib  # type: ignore
 
 
 @dataclass(frozen=True)
@@ -281,7 +281,7 @@ def in_colab():
     if in_colab():
         return True
     try:
-        shell = get_ipython().__class__.__name__
+        shell = get_ipython().__class__.__name__  # type: ignore
         if shell == "ZMQInteractiveShell":  # Jupyter notebook, Spyder or qtconsole
             try:
                 from IPython.display import (
@@ -398,6 +398,8 @@ def _extract(self, results):
         dims["divergence_start_gradient"] = ["unconstrained_parameter"]
         dims["divergence_end"] = ["unconstrained_parameter"]
         dims["divergence_momentum"] = ["unconstrained_parameter"]
+        dims["transformed_gradient"] = ["unconstrained_parameter"]
+        dims["transformed_position"] = ["unconstrained_parameter"]
 
         if self._return_raw_trace:
             return results
@@ -453,14 +455,15 @@ def _repr_html_(self):
 def sample(
     compiled_model: CompiledModel,
     *,
-    draws: int,
-    tune: int,
+    draws: int | None,
+    tune: int | None,
     chains: int,
     cores: Optional[int],
     seed: Optional[int],
     save_warmup: bool,
     progress_bar: bool,
     low_rank_modified_mass_matrix: bool = False,
+    transform_adapt: bool = False,
     init_mean: Optional[np.ndarray],
     return_raw_trace: bool,
     blocking: Literal[True],
@@ -472,14 +475,15 @@ def sample(
 def sample(
     compiled_model: CompiledModel,
     *,
-    draws: int,
-    tune: int,
+    draws: int | None,
+    tune: int | None,
     chains: int,
     cores: Optional[int],
     seed: Optional[int],
     save_warmup: bool,
     progress_bar: bool,
     low_rank_modified_mass_matrix: bool = False,
+    transform_adapt: bool = False,
     init_mean: Optional[np.ndarray],
     return_raw_trace: bool,
     blocking: Literal[False],
@@ -490,14 +494,15 @@ def sample(
 def sample(
     compiled_model: CompiledModel,
     *,
-    draws: int = 1000,
-    tune: int = 300,
+    draws: int | None = None,
+    tune: int | None = None,
     chains: int = 6,
     cores: Optional[int] = None,
     seed: Optional[int] = None,
     save_warmup: bool = True,
     progress_bar: bool = True,
     low_rank_modified_mass_matrix: bool = False,
+    transform_adapt: bool = False,
     init_mean: Optional[np.ndarray] = None,
     return_raw_trace: bool = False,
     blocking: bool = True,
@@ -510,9 +515,9 @@ def sample(
 
     Parameters
     ----------
-    draws: int
+    draws: int | None
         The number of draws after tuning in each chain.
-    tune: int
+    tune: int | None
         The number of tuning (warmup) draws in each chain.
     chains: int
         The number of chains to sample.
@@ -585,6 +590,9 @@ def sample(
     mass_matrix_gamma: float > 0, default=1e-5
         Regularisation parameter for the eigenvalues. Only
         applicable with low_rank_modified_mass_matrix=True.
+    transform_adapt: bool, default=False
+        Use the experimental transform adaptation algorithm
+        during tuning.
     **kwargs
         Pass additional arguments to nutpie._lib.PySamplerArgs
 
@@ -594,12 +602,22 @@ def sample(
         An ArviZ ``InferenceData`` object that contains the samples.
     """
 
+    if low_rank_modified_mass_matrix and transform_adapt:
+        raise ValueError(
+            "Specify only one of `low_rank_modified_mass_matrix` and `transform_adapt`"
+        )
+
     if low_rank_modified_mass_matrix:
         settings = _lib.PyNutsSettings.LowRank(seed)
+    elif transform_adapt:
+        settings = _lib.PyNutsSettings.Transform(seed)
     else:
         settings = _lib.PyNutsSettings.Diag(seed)
-    settings.num_tune = tune
-    settings.num_draws = draws
+
+    if tune is not None:
+        settings.num_tune = tune
+    if draws is not None:
+        settings.num_draws = draws
     settings.num_chains = chains
 
     for name, val in kwargs.items():
@@ -608,10 +626,10 @@ def sample(
     if cores is None:
         try:
             # Only available in python>=3.13
-            available = os.process_cpu_count()
+            available = os.process_cpu_count()  # type: ignore
         except AttributeError:
             available = os.cpu_count()
-        cores = min(chains, available)
+        cores = min(chains, cast(int, available))
 
     if init_mean is None:
         init_mean = np.zeros(compiled_model.n_dim)
diff --git a/python/nutpie/transform_adapter.py b/python/nutpie/transform_adapter.py
new file mode 100644
index 0000000..9a05054
--- /dev/null
+++ b/python/nutpie/transform_adapter.py
@@ -0,0 +1,627 @@
+from functools import partial
+import math
+
+import numpy as np
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import traceback
+import flowjax
+import flowjax.flows
+import flowjax.train
+import optax
+from paramax import unwrap
+
+from nutpie.normalizing_flow import extend_flow, make_flow
+
+_BIJECTION_TRACE = []
+
+
+class FisherLoss:
+    def __init__(self, gamma=None, log_inside_batch=False):
+        self._gamma = gamma
+        self._log_inside_batch = log_inside_batch
+
+    @eqx.filter_jit
+    def __call__(
+        self,
+        params,
+        static,
+        draws,
+        grads,
+        logps,
+        condition=None,
+        key=None,
+        return_all_costs=False,
+        return_elemwise_costs=False,
+    ):
+        flow = unwrap(eqx.combine(params, static, is_leaf=eqx.is_inexact_array))
+
+        if return_elemwise_costs:
+
+            def compute_loss(bijection, draw, grad, logp):
+                if True:
+                    draw, grad, logp = bijection.inverse_gradient_and_val_(
+                        draw, grad, logp
+                    )
+                else:
+                    draw, grad, logp = (
+                        flowjax.bijections.AbstractBijection.inverse_gradient_and_val_(
+                            bijection, draw, grad, logp
+                        )
+                    )
+                cost = (draw + grad) ** 2
+                return cost
+
+            costs = jax.vmap(compute_loss, [None, 0, 0, 0])(
+                flow,
+                draws,
+                grads,
+                logps,
+            )
+            return costs.mean(0)
+
+        if self._gamma is None:
+
+            def compute_loss(bijection, draw, grad, logp):
+                draw, grad, logp = bijection.inverse_gradient_and_val_(draw, grad, logp)
+                cost = ((draw + grad) ** 2).sum()
+                return cost
+
+            costs = jax.vmap(compute_loss, [None, 0, 0, 0])(
+                flow,
+                draws,
+                grads,
+                logps,
+            )
+
+            if return_all_costs:
+                return costs
+
+            if self._log_inside_batch:
+                return jnp.log(costs).mean()
+            else:
+                return jnp.log(costs.mean())
+
+        else:
+
+            def transform(draw, grad, logp):
+                return flow.inverse_gradient_and_val_(draw, grad, logp)
+
+            draws, grads, logps = jax.vmap(transform, [0, 0, 0], (0, 0, 0))(
+                draws, grads, logps
+            )
+            fisher_loss = ((draws + grads) ** 2).sum(1).mean(0)
+            normal_logps = -(draws * draws).sum(1) / 2
+            var_loss = (logps - normal_logps).var()
+            return jnp.log(fisher_loss + self._gamma * var_loss)
+
+
+def fit_flow(key, bijection, loss_fn, draws, grads, logps, **kwargs):
+    flow = bijection
+
+    key, train_key = jax.random.split(key)
+
+    fit, losses = flowjax.train.fit_to_data(
+        key=train_key,
+        dist=flow,
+        x=(draws, grads, logps),
+        loss_fn=loss_fn,
+        max_epochs=1000,
+        return_best=True,
+        **kwargs,
+    )
+    return fit, losses, losses["opt_state"]
+
+
+@eqx.filter_jit
+def _init_from_transformed_position(logp_fn, bijection, transformed_position):
+    bijection = unwrap(bijection)
+    (untransformed_position, logdet), pull_grad = jax.vjp(
+        bijection.transform_and_log_det, transformed_position
+    )
+    logp, untransformed_gradient = jax.value_and_grad(lambda x: logp_fn(x)[0])(
+        untransformed_position
+    )
+    (transformed_gradient,) = pull_grad((untransformed_gradient, 1.0))
+    return (
+        logp,
+        logdet,
+        untransformed_position,
+        untransformed_gradient,
+        transformed_gradient,
+    )
+
+
+@eqx.filter_jit
+def _init_from_transformed_position_part1(logp_fn, bijection, transformed_position):
+    bijection = unwrap(bijection)
+    (untransformed_position, logdet) = bijection.transform_and_log_det(
+        transformed_position
+    )
+
+    return (logdet, untransformed_position)
+
+
+@eqx.filter_jit
+def _init_from_transformed_position_part2(
+    bijection,
+    part1,
+    untransformed_gradient,
+):
+    logdet, untransformed_position, transformed_position = part1
+    bijection = unwrap(bijection)
+    _, pull_grad = jax.vjp(bijection.transform_and_log_det, transformed_position)
+    (transformed_gradient,) = pull_grad((untransformed_gradient, 1.0))
+    return (
+        logdet,
+        transformed_gradient,
+    )
+
+
+@eqx.filter_jit
+def _init_from_untransformed_position(logp_fn, bijection, untransformed_position):
+    logp, untransformed_gradient = jax.value_and_grad(lambda x: logp_fn(x)[0])(
+        untransformed_position
+    )
+    logdet, transformed_position, transformed_gradient = _inv_transform(
+        bijection, untransformed_position, untransformed_gradient
+    )
+    return (
+        logp,
+        logdet,
+        untransformed_gradient,
+        transformed_position,
+        transformed_gradient,
+    )
+
+
+@eqx.filter_jit
+def _inv_transform(bijection, untransformed_position, untransformed_gradient):
+    bijection = unwrap(bijection)
+    transformed_position, transformed_gradient, logdet = (
+        bijection.inverse_gradient_and_val_(
+            untransformed_position, untransformed_gradient, 0.0
+        )
+    )
+    return logdet, transformed_position, transformed_gradient
+
+
+class TransformAdapter:
+    def __init__(
+        self,
+        seed,
+        position,
+        gradient,
+        chain,
+        *,
+        logp_fn,
+        make_flow_fn,
+        verbose=False,
+        window_size=2000,
+        show_progress=False,
+        num_diag_windows=10,
+        learning_rate=1e-3,
+        zero_init=True,
+        untransformed_dim=None,
+        batch_size=128,
+        reuse_opt_state=True,
+        max_patience=5,
+        gamma=None,
+        log_inside_batch=False,
+        initial_skip=500,
+        extension_windows=None,
+        extend_dct=False,
+        extension_var_count=6,
+        extension_var_trafo_count=4,
+        debug_save_bijection=False,
+        make_optimizer=None,
+        num_layers=9,
+    ):
+        self._logp_fn = logp_fn
+        self._make_flow_fn = make_flow_fn
+        self._chain = chain
+        self._verbose = verbose
+        self._window_size = window_size
+        self._initial_skip = initial_skip
+        self._num_layers = num_layers
+        if make_optimizer is None:
+            self._make_optimizer = lambda: optax.apply_if_finite(
+                optax.adamw(learning_rate), 50
+            )
+        else:
+            self._make_optimizer = make_optimizer
+        self._optimizer = self._make_optimizer()
+        self._loss_fn = FisherLoss(gamma, log_inside_batch)
+        self._show_progress = show_progress
+        self._num_diag_windows = num_diag_windows
+        self._zero_init = zero_init
+        self._untransformed_dim = untransformed_dim
+        self._batch_size = batch_size
+        self._reuse_opt_state = reuse_opt_state
+        self._opt_state = None
+        self._max_patience = max_patience
+        self._count_trace = []
+        self._last_extend_dct = True
+        self._extend_dct = extend_dct
+        self._extension_var_count = extension_var_count
+        self._extension_var_trafo_count = extension_var_trafo_count
+        self._debug_save_bijection = debug_save_bijection
+        self._layers = 0
+
+        if extension_windows is None:
+            self._extension_windows = []
+        else:
+            self._extension_windows = extension_windows
+
+        try:
+            self._bijection = make_flow_fn(seed, [position], [gradient], n_layers=0)
+        except Exception as e:
+            print("make_flow", e)
+            print(traceback.format_exc())
+            raise
+        self.index = 0
+
+    @property
+    def transformation_id(self):
+        return self.index
+
+    def update(self, seed, positions, gradients, logps):
+        self.index += 1
+        if self._verbose:
+            print(f"Chain {self._chain}: Total available points: {len(positions)}")
+        n_draws = len(positions)
+        assert n_draws == len(positions)
+        assert n_draws == len(gradients)
+        assert n_draws == len(logps)
+        self._count_trace.append(n_draws)
+        if n_draws == 0:
+            return
+        try:
+            if self.index <= self._num_diag_windows:
+                size = len(positions)
+                lower_idx = -size // 5 + 3
+                positions_slice = positions[lower_idx:]
+                gradients_slice = gradients[lower_idx:]
+                logp_slice = logps[lower_idx:]
+
+                if len(positions_slice) > 0:
+                    positions = positions_slice
+                    gradients = gradients_slice
+                    logps = logp_slice
+
+                positions = np.array(positions)
+                gradients = np.array(gradients)
+                logps = np.array(logps)
+
+                fit = self._make_flow_fn(seed, positions, gradients, n_layers=0)
+
+                flow = fit
+                params, static = eqx.partition(flow, eqx.is_inexact_array)
+                new_loss = self._loss_fn(params, static, positions, gradients, logps)
+
+                if self._verbose:
+                    print("loss from diag:", new_loss)
+
+                if np.isfinite(new_loss):
+                    self._bijection = fit
+                    self._opt_state = None
+
+                return
+
+            positions = np.array(positions[self._initial_skip :][-self._window_size :])
+            gradients = np.array(gradients[self._initial_skip :][-self._window_size :])
+            logps = np.array(logps[self._initial_skip :][-self._window_size :])
+
+            if len(positions) < 10:
+                return
+
+            if self._verbose and not np.isfinite(gradients).all():
+                print(gradients)
+                print(gradients.shape)
+                print((~np.isfinite(gradients)).nonzero())
+
+            assert np.isfinite(positions).all()
+            assert np.isfinite(gradients).all()
+            assert np.isfinite(logps).all()
+
+            # TODO don't reuse seed
+            key = jax.random.PRNGKey(seed % (2**63))
+
+            if len(self._bijection.bijections) == 1:
+                base = self._make_flow_fn(
+                    seed,
+                    positions,
+                    gradients,
+                    n_layers=self._num_layers,
+                    untransformed_dim=self._untransformed_dim,
+                    zero_init=self._zero_init,
+                )
+                flow = base
+                params, static = eqx.partition(flow, eqx.is_inexact_array)
+                if self._verbose:
+                    print(
+                        "number of parameters in flow:",
+                        sum([math.prod(val.shape) for val in jax.tree.leaves(params)]),
+                    )
+
+                    print(
+                        "loss before optimization: ",
+                        self._loss_fn(
+                            params,
+                            static,
+                            positions[-128:],
+                            gradients[-128:],
+                            logps[-128:],
+                        ),
+                    )
+            else:
+                base = self._bijection
+
+            if self.index in self._extension_windows:
+                if self._verbose:
+                    print("Extending flow...")
+                self._last_extend_dct = not self._last_extend_dct
+                dct = self._last_extend_dct and self._extend_dct
+                base = extend_flow(
+                    key,
+                    base,
+                    self._loss_fn,
+                    positions,
+                    gradients,
+                    logps,
+                    self._layers,
+                    dct=dct,
+                    extension_var_count=self._extension_var_count,
+                    extension_var_trafo_count=self._extension_var_trafo_count,
+                    verbose=self._verbose,
+                )
+                self._optimizer = self._make_optimizer()
+                self._opt_state = None
+                self._layers += 1
+
+            # make_flow might still onreturn a single trafo for 1d problems
+            if len(base.bijections) == 1:
+                self._bijection = base
+                self._opt_state = None
+                return
+
+            flow = self._bijection
+            params, static = eqx.partition(flow, eqx.is_inexact_array)
+            old_loss = self._loss_fn(
+                params, static, positions[-128:], gradients[-128:], logps[-128:]
+            )
+
+            if np.isfinite(old_loss) and old_loss < -5 and self.index > 10:
+                if self._verbose:
+                    print(f"Loss is low ({old_loss}), skipping training")
+                return
+
+            fit, _, opt_state = fit_flow(
+                key,
+                base,
+                self._loss_fn,
+                positions,
+                gradients,
+                logps,
+                show_progress=self._show_progress,
+                optimizer=self._optimizer,
+                batch_size=self._batch_size,
+                opt_state=self._opt_state if self._reuse_opt_state else None,
+                max_patience=self._max_patience,
+            )
+
+            flow = fit
+            params, static = eqx.partition(flow, eqx.is_inexact_array)
+            new_loss = self._loss_fn(
+                params, static, positions[-128:], gradients[-128:], logps[-128:]
+            )
+
+            if self._verbose:
+                print(f"Chain {self._chain}: New loss {new_loss}, old loss {old_loss}")
+
+            if not np.isfinite(old_loss):
+                flow = self._bijection
+                params, static = eqx.partition(flow, eqx.is_inexact_array)
+                print(
+                    self._loss_fn(
+                        params,
+                        static,
+                        positions[-128:],
+                        gradients[-128:],
+                        logps[-128:],
+                        return_all_costs=True,
+                    )
+                )
+
+            if not np.isfinite(new_loss):
+                flow = fit
+                params, static = eqx.partition(flow, eqx.is_inexact_array)
+                print(
+                    self._loss_fn(
+                        params,
+                        static,
+                        positions[-128:],
+                        gradients[-128:],
+                        logps[-128:],
+                        return_all_costs=True,
+                    )
+                )
+
+            if self._debug_save_bijection:
+                _BIJECTION_TRACE.append(
+                    (self.index, fit, (positions, gradients, logps))
+                )
+
+            def valid_new_logp():
+                logdet, pos, grad = _inv_transform(
+                    fit,
+                    jnp.array(positions[-1]),
+                    jnp.array(gradients[-1]),
+                )
+                return (
+                    np.isfinite(logdet)
+                    and np.isfinite(pos[0]).all()
+                    and np.isfinite(grad[0]).all()
+                )
+
+            if (not np.isfinite(old_loss)) and (not np.isfinite(new_loss)):
+                self._bijection = self._make_flow_fn(
+                    seed, positions, gradients, n_layers=0
+                )
+                self._opt_state = None
+                return
+
+            if not valid_new_logp():
+                if self._verbose:
+                    print("Invalid new logp. Skipping update.")
+                return
+
+            if not np.isfinite(new_loss):
+                if self._verbose:
+                    print("Invalid new loss. Skipping update.")
+                return
+
+            if new_loss > old_loss:
+                return
+
+            self._bijection = fit
+            self._opt_state = opt_state
+
+        except Exception as e:
+            print("update error:", e)
+            print(traceback.format_exc())
+            raise
+
+    def init_from_transformed_position(self, transformed_position):
+        try:
+            logp, logdet, *arrays = _init_from_transformed_position(
+                self._logp_fn,
+                self._bijection,
+                jnp.array(transformed_position),
+            )
+            return (
+                float(logp),
+                float(logdet),
+                *[np.array(val, dtype="float64") for val in arrays],
+            )
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+            raise
+
+    def init_from_transformed_position_part1(self, transformed_position):
+        try:
+            transformed_position = jnp.array(transformed_position)
+            logdet, untransformed_position = _init_from_transformed_position_part1(
+                self._logp_fn,
+                self._bijection,
+                transformed_position,
+            )
+            part1 = (logdet, untransformed_position, transformed_position)
+            return np.array(untransformed_position, dtype="float64"), part1
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+            raise
+
+    def init_from_transformed_position_part2(
+        self,
+        part1,
+        untransformed_gradient,
+    ):
+        try:
+            # TODO We could extract the arrays from the pull_grad function
+            # to reuse computation from part1
+            logdet, *arrays = _init_from_transformed_position_part2(
+                self._bijection,
+                part1,
+                untransformed_gradient,
+            )
+            return float(logdet), *[np.array(val, dtype="float64") for val in arrays]
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+            raise
+
+    def init_from_untransformed_position(self, untransformed_position):
+        try:
+            logp, logdet, *arrays = _init_from_untransformed_position(
+                self._logp_fn,
+                self._bijection,
+                jnp.array(untransformed_position),
+            )
+            arrays = [np.array(val, dtype="float64") for val in arrays]
+            return float(logp), float(logdet), *arrays
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+            raise
+
+    def inv_transform(self, position, gradient):
+        try:
+            logdet, *arrays = _inv_transform(
+                self._bijection, jnp.array(position), jnp.array(gradient)
+            )
+            return logdet, *[np.array(val, dtype="float64") for val in arrays]
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+            raise
+
+
+def make_transform_adapter(
+    *,
+    verbose=False,
+    window_size=600,
+    show_progress=False,
+    nn_depth=1,
+    nn_width=None,
+    num_layers=9,
+    num_diag_windows=9,
+    learning_rate=5e-4,
+    untransformed_dim=None,
+    zero_init=True,
+    batch_size=128,
+    reuse_opt_state=False,
+    max_patience=20,
+    householder_layer=False,
+    dct_layer=False,
+    gamma=None,
+    log_inside_batch=False,
+    initial_skip=120,
+    extension_windows=[],
+    extend_dct=False,
+    extension_var_count=4,
+    extension_var_trafo_count=2,
+    debug_save_bijection=False,
+    make_optimizer=None,
+):
+    return partial(
+        TransformAdapter,
+        verbose=verbose,
+        window_size=window_size,
+        make_flow_fn=partial(
+            make_flow,
+            householder_layer=householder_layer,
+            dct_layer=dct_layer,
+            nn_width=nn_width,
+        ),
+        show_progress=show_progress,
+        num_diag_windows=num_diag_windows,
+        learning_rate=learning_rate,
+        zero_init=zero_init,
+        untransformed_dim=untransformed_dim,
+        batch_size=batch_size,
+        reuse_opt_state=reuse_opt_state,
+        max_patience=max_patience,
+        gamma=gamma,
+        log_inside_batch=log_inside_batch,
+        initial_skip=initial_skip,
+        extension_windows=extension_windows,
+        extend_dct=extend_dct,
+        extension_var_count=extension_var_count,
+        extension_var_trafo_count=extension_var_trafo_count,
+        debug_save_bijection=debug_save_bijection,
+        make_optimizer=make_optimizer,
+    )
diff --git a/src/progress.rs b/src/progress.rs
index 906c50d..2403e15 100644
--- a/src/progress.rs
+++ b/src/progress.rs
@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, time::Duration};
+use std::{collections::BTreeMap, sync::Arc, time::Duration};
 
 use anyhow::{Context, Result};
 use indicatif::ProgressBar;
@@ -10,13 +10,13 @@ use upon::{Engine, Value};
 pub struct ProgressHandler {
     engine: Engine<'static>,
     template: String,
-    callback: Py<PyAny>,
+    callback: Arc<Py<PyAny>>,
     rate: Duration,
     n_cores: usize,
 }
 
 impl ProgressHandler {
-    pub fn new(callback: Py<PyAny>, rate: Duration, template: String, n_cores: usize) -> Self {
+    pub fn new(callback: Arc<Py<PyAny>>, rate: Duration, template: String, n_cores: usize) -> Self {
         let engine = Engine::new();
         Self {
             engine,
diff --git a/src/pyfunc.rs b/src/pyfunc.rs
index f914974..f23145e 100644
--- a/src/pyfunc.rs
+++ b/src/pyfunc.rs
@@ -8,9 +8,10 @@ use arrow::{
     },
     datatypes::{DataType, Field, Float32Type, Float64Type, Int64Type},
 };
-use numpy::{PyArray1, PyReadonlyArray1};
+use numpy::{NotContiguousError, PyArray1, PyReadonlyArray1};
 use nuts_rs::{CpuLogpFunc, CpuMath, DrawStorage, LogpError, Model};
 use pyo3::{
+    exceptions::PyRuntimeError,
     pyclass, pymethods,
     types::{PyAnyMethods, PyDict, PyDictMethods},
     Bound, Py, PyAny, PyErr, Python,
@@ -20,6 +21,8 @@ use rand_distr::{Distribution, Uniform};
 use smallvec::SmallVec;
 use thiserror::Error;
 
+use crate::wrapper::PyTransformAdapt;
+
 #[pyclass]
 #[derive(Debug, Clone)]
 #[non_exhaustive]
@@ -71,29 +74,33 @@ impl PyVariable {
 #[pyclass]
 #[derive(Debug, Clone)]
 pub struct PyModel {
-    make_logp_func: Py<PyAny>,
-    make_expand_func: Py<PyAny>,
-    init_point_func: Option<Py<PyAny>>,
-    variables: Vec<PyVariable>,
+    make_logp_func: Arc<Py<PyAny>>,
+    make_expand_func: Arc<Py<PyAny>>,
+    init_point_func: Option<Arc<Py<PyAny>>>,
+    variables: Arc<Vec<PyVariable>>,
+    transform_adapter: Option<PyTransformAdapt>,
     ndim: usize,
 }
 
 #[pymethods]
 impl PyModel {
     #[new]
+    #[pyo3(signature = (make_logp_func, make_expand_func, variables, ndim, *, init_point_func=None, transform_adapter=None))]
     fn new<'py>(
         make_logp_func: Py<PyAny>,
         make_expand_func: Py<PyAny>,
         variables: Vec<PyVariable>,
         ndim: usize,
         init_point_func: Option<Py<PyAny>>,
+        transform_adapter: Option<Py<PyAny>>,
     ) -> Self {
         Self {
-            make_logp_func,
-            make_expand_func,
-            init_point_func,
-            variables,
+            make_logp_func: Arc::new(make_logp_func),
+            make_expand_func: Arc::new(make_expand_func),
+            init_point_func: init_point_func.map(|x| x.into()),
+            variables: Arc::new(variables),
             ndim,
+            transform_adapter: transform_adapter.map(PyTransformAdapt::new),
         }
     }
 }
@@ -103,9 +110,13 @@ pub enum PyLogpError {
     #[error("Bad logp value: {0}")]
     BadLogp(f64),
     #[error("Python error: {0}")]
-    PyError(PyErr),
+    PyError(#[from] PyErr),
     #[error("logp function must return float.")]
     ReturnTypeError(),
+    #[error("Python retured a non-contigous array")]
+    NotContiguousError(#[from] NotContiguousError),
+    #[error("Unknown error: {0}")]
+    Anyhow(#[from] anyhow::Error),
 }
 
 impl LogpError for PyLogpError {
@@ -113,7 +124,7 @@ impl LogpError for PyLogpError {
         match self {
             Self::BadLogp(_) => true,
             Self::PyError(err) => Python::with_gil(|py| {
-                let Ok(attr) = err.value_bound(py).getattr("is_recoverable") else {
+                let Ok(attr) = err.value(py).getattr("is_recoverable") else {
                     return false;
                 };
                 return attr
@@ -121,20 +132,29 @@ impl LogpError for PyLogpError {
                     .expect("Could not access is_recoverable in error check");
             }),
             Self::ReturnTypeError() => false,
+            Self::NotContiguousError(_) => false,
+            Self::Anyhow(_) => false,
         }
     }
 }
 
 pub struct PyDensity {
     logp: Py<PyAny>,
+    transform_adapter: Option<PyTransformAdapt>,
     dim: usize,
 }
 
 impl PyDensity {
-    fn new(logp_clone_func: &Py<PyAny>, dim: usize) -> Result<Self> {
+    fn new(
+        logp_clone_func: &Py<PyAny>,
+        dim: usize,
+        transform_adapter: Option<&PyTransformAdapt>,
+    ) -> Result<Self> {
         let logp_func = Python::with_gil(|py| logp_clone_func.call0(py))?;
+        let transform_adapter = transform_adapter.map(|val| val.clone());
         Ok(Self {
             logp: logp_func,
+            transform_adapter,
             dim,
         })
     }
@@ -142,10 +162,11 @@ impl PyDensity {
 
 impl CpuLogpFunc for PyDensity {
     type LogpError = PyLogpError;
+    type TransformParams = Py<PyAny>;
 
     fn logp(&mut self, position: &[f64], grad: &mut [f64]) -> Result<f64, Self::LogpError> {
         Python::with_gil(|py| {
-            let pos_array = PyArray1::from_slice_bound(py, position);
+            let pos_array = PyArray1::from_slice(py, position);
             let result = self.logp.call1(py, (pos_array,));
             match result {
                 Ok(val) => {
@@ -172,11 +193,122 @@ impl CpuLogpFunc for PyDensity {
     fn dim(&self) -> usize {
         self.dim
     }
+
+    fn inv_transform_normalize(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<f64, Self::LogpError> {
+        let logdet = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .inv_transform_normalize(
+                params,
+                untransformed_position,
+                untransformed_gradient,
+                transformed_position,
+                transformed_gradient,
+            )?;
+        Ok(logdet)
+    }
+
+    fn init_from_transformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &mut [f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &[f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<(f64, f64), Self::LogpError> {
+        let (logp, logdet) = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .init_from_transformed_position(
+                params,
+                untransformed_position,
+                untransformed_gradient,
+                transformed_position,
+                transformed_gradient,
+            )?;
+        Ok((logp, logdet))
+    }
+
+    fn init_from_untransformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<(f64, f64), Self::LogpError> {
+        let (logp, logdet) = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .init_from_untransformed_position(
+                params,
+                untransformed_position,
+                untransformed_gradient,
+                transformed_position,
+                transformed_gradient,
+            )?;
+        Ok((logp, logdet))
+    }
+
+    fn update_transformation<'a, R: rand::Rng + ?Sized>(
+        &'a mut self,
+        rng: &mut R,
+        untransformed_positions: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_gradients: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_logp: impl ExactSizeIterator<Item = &'a f64>,
+        params: &'a mut Py<PyAny>,
+    ) -> std::result::Result<(), Self::LogpError> {
+        self.transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .update_transformation(
+                rng,
+                untransformed_positions,
+                untransformed_gradients,
+                untransformed_logp,
+                params,
+            )?;
+        Ok(())
+    }
+
+    fn new_transformation<R: rand::Rng + ?Sized>(
+        &mut self,
+        rng: &mut R,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        chain: u64,
+    ) -> std::result::Result<Py<PyAny>, Self::LogpError> {
+        let trafo = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .new_transformation(rng, untransformed_position, untransformed_gradient, chain)?;
+        Ok(trafo)
+    }
+
+    fn transformation_id(&self, params: &Py<PyAny>) -> std::result::Result<i64, Self::LogpError> {
+        let id = self
+            .transform_adapter
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .transformation_id(params)?;
+        Ok(id)
+    }
 }
 
 pub struct PyTrace {
     expand: Py<PyAny>,
-    variables: Vec<PyVariable>,
+    variables: Arc<Vec<PyVariable>>,
     builder: StructBuilder,
 }
 
@@ -184,7 +316,7 @@ impl PyTrace {
     pub fn new<R: Rng + ?Sized>(
         rng: &mut R,
         chain: u64,
-        variables: Vec<PyVariable>,
+        variables: Arc<Vec<PyVariable>>,
         make_expand_func: &Py<PyAny>,
         capacity: usize,
     ) -> Result<Self> {
@@ -234,6 +366,7 @@ impl TensorShape {
 #[pymethods]
 impl TensorShape {
     #[new]
+    #[pyo3(signature = (shape, dims=None))]
     fn py_new(shape: Vec<usize>, dims: Option<Vec<Option<String>>>) -> Result<Self> {
         let dims = dims.unwrap_or(shape.iter().map(|_| None).collect());
         if dims.len() != shape.len() {
@@ -318,7 +451,7 @@ impl ExpandDtype {
 impl DrawStorage for PyTrace {
     fn append_value(&mut self, point: &[f64]) -> Result<()> {
         Python::with_gil(|py| {
-            let point = PyArray1::from_slice_bound(py, point);
+            let point = PyArray1::from_slice(py, point);
             let full_point = self
                 .expand
                 .call1(py, (point,))
@@ -344,36 +477,51 @@ impl DrawStorage for PyTrace {
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type",
                                 )?;
-                            builder.append_value(value.extract().expect("Return value from expand function could not be converted to boolean"))
+                            let value = value
+                                .extract()
+                                .expect("Return value from expand function could not be converted to boolean");
+                            builder.append_value(value)
                         },
                         ExpandDtype::Float64 {} => {
                             let builder: &mut Float64Builder =
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type",
                                 )?;
-                            builder.append_value(value.extract().expect("Return value from expand function could not be converted to float64"))
+                            builder.append_value(
+                                value
+                                .extract()
+                                .expect("Return value from expand function could not be converted to float64")
+                            )
                         },
                         ExpandDtype::Float32 {} => {
                             let builder: &mut Float32Builder =
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type",
                                 )?;
-                            builder.append_value(value.extract().expect("Return value from expand function could not be converted to float32"))
-
+                            builder.append_value(
+                                value
+                                .extract()
+                                .expect("Return value from expand function could not be converted to float32")
+                            )
                         },
                         ExpandDtype::Int64 {} => {
                             let builder: &mut Int64Builder =
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type",
                                 )?;
-                            builder.append_value(value.extract().expect("Return value from expand function could not be converted to int64"))
+                            let value = value.extract().expect("Return value from expand function could not be converted to int64");
+                            builder.append_value(value)
                         },
                         ExpandDtype::BooleanArray { tensor_type } => {
                             let builder: &mut LargeListBuilder<Box<dyn ArrayBuilder>> =
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type. Expected LargeListBuilder of Bool",
                                 )?;
-                            let value_builder = builder.values().as_any_mut().downcast_mut::<BooleanBuilder>().context("Could not downcast builder to boolean type")?;
+                            let value_builder = builder
+                                .values()
+                                .as_any_mut()
+                                .downcast_mut::<BooleanBuilder>()
+                                .context("Could not downcast builder to boolean type")?;
                             let values: PyReadonlyArray1<bool> = value.extract().context("Could not convert object to array")?;
                             if values.len()? != tensor_type.size() {
                                 bail!("Extracted array has incorrect shape");
@@ -387,7 +535,11 @@ impl DrawStorage for PyTrace {
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type. Expected LargeListBuilder of Float64",
                                 )?;
-                            let value_builder = builder.values().as_any_mut().downcast_mut::<PrimitiveBuilder<Float64Type>>().context("Could not downcast builder to float64 type")?;
+                            let value_builder = builder
+                                .values()
+                                .as_any_mut()
+                                .downcast_mut::<PrimitiveBuilder<Float64Type>>()
+                                .context("Could not downcast builder to float64 type")?;
                             let values: PyReadonlyArray1<f64> = value.extract().context("Could not convert object to array")?;
                             if values.len()? != tensor_type.size() {
                                 bail!("Extracted array has incorrect shape");
@@ -400,7 +552,11 @@ impl DrawStorage for PyTrace {
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type. Expected LargeListBuilder of Float32",
                                 )?;
-                            let value_builder = builder.values().as_any_mut().downcast_mut::<PrimitiveBuilder<Float32Type>>().context("Could not downcast builder to float32 type")?;
+                            let value_builder = builder
+                                .values()
+                                .as_any_mut()
+                                .downcast_mut::<PrimitiveBuilder<Float32Type>>()
+                                .context("Could not downcast builder to float32 type")?;
                             let values: PyReadonlyArray1<f32> = value.extract().context("Could not convert object to array")?;
                             if values.len()? != tensor_type.size() {
                                 bail!("Extracted array has incorrect shape");
@@ -413,7 +569,11 @@ impl DrawStorage for PyTrace {
                                 self.builder.field_builder(i).context(
                                     "Builder has incorrect type. Expected LargeListBuilder of Int64",
                                 )?;
-                            let value_builder = builder.values().as_any_mut().downcast_mut::<PrimitiveBuilder<Int64Type>>().context("Could not downcast builder to i64 type")?;
+                            let value_builder = builder
+                                .values()
+                                .as_any_mut()
+                                .downcast_mut::<PrimitiveBuilder<Int64Type>>()
+                                .context("Could not downcast builder to i64 type")?;
                             let values: PyReadonlyArray1<i64> = value.extract().context("Could not convert object to array")?;
                             if values.len()? != tensor_type.size() {
                                 bail!("Extracted array has incorrect shape");
@@ -471,6 +631,7 @@ impl Model for PyModel {
         Ok(CpuMath::new(PyDensity::new(
             &self.make_logp_func,
             self.ndim,
+            self.transform_adapter.as_ref(),
         )?))
     }
 
@@ -480,7 +641,7 @@ impl Model for PyModel {
         position: &mut [f64],
     ) -> Result<()> {
         let Some(init_func) = self.init_point_func.as_ref() else {
-            let dist = Uniform::new(-2f64, 2f64);
+            let dist = Uniform::new(-2f64, 2f64).expect("Could not create uniform distribution");
             position.iter_mut().for_each(|x| *x = dist.sample(rng));
             return Ok(());
         };
diff --git a/src/pymc.rs b/src/pymc.rs
index 98426c7..b33b821 100644
--- a/src/pymc.rs
+++ b/src/pymc.rs
@@ -39,7 +39,7 @@ type RawExpandFunc = unsafe extern "C" fn(
 #[derive(Clone)]
 pub(crate) struct LogpFunc {
     func: RawLogpFunc,
-    _keep_alive: PyObject,
+    _keep_alive: Arc<PyObject>,
     user_data_ptr: UserData,
     dim: usize,
 }
@@ -55,7 +55,7 @@ impl LogpFunc {
             unsafe { std::mem::transmute::<*const c_void, RawLogpFunc>(ptr as *const c_void) };
         Self {
             func,
-            _keep_alive: keep_alive,
+            _keep_alive: Arc::new(keep_alive),
             user_data_ptr: user_data_ptr as UserData,
             dim,
         }
@@ -66,7 +66,7 @@ impl LogpFunc {
 #[derive(Clone)]
 pub(crate) struct ExpandFunc {
     func: RawExpandFunc,
-    _keep_alive: PyObject,
+    _keep_alive: Arc<PyObject>,
     user_data_ptr: UserData,
     dim: usize,
     expanded_dim: usize,
@@ -87,7 +87,7 @@ impl ExpandFunc {
         Self {
             dim,
             expanded_dim,
-            _keep_alive: keep_alive,
+            _keep_alive: Arc::new(keep_alive),
             user_data_ptr: user_data_ptr as UserData,
             func,
         }
@@ -114,6 +114,7 @@ impl LogpError for ErrorCode {
 
 impl<'a> CpuLogpFunc for &'a LogpFunc {
     type LogpError = ErrorCode;
+    type TransformParams = ();
 
     fn dim(&self) -> usize {
         self.dim
@@ -233,7 +234,7 @@ pub(crate) struct PyMcModel {
     dim: usize,
     density: LogpFunc,
     expand: ExpandFunc,
-    init_func: Py<PyAny>,
+    init_func: Arc<Py<PyAny>>,
     var_sizes: Vec<usize>,
     var_names: Vec<String>,
 }
@@ -253,7 +254,7 @@ impl PyMcModel {
             dim,
             density,
             expand,
-            init_func,
+            init_func: init_func.into(),
             var_names: var_names.extract()?,
             var_sizes: var_sizes.extract()?,
         })
diff --git a/src/stan.rs b/src/stan.rs
index 0564cd2..e258096 100644
--- a/src/stan.rs
+++ b/src/stan.rs
@@ -7,16 +7,19 @@ use arrow::datatypes::{DataType, Field};
 use bridgestan::open_library;
 use itertools::{izip, Itertools};
 use nuts_rs::{CpuLogpFunc, CpuMath, DrawStorage, LogpError, Model, Settings};
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 use pyo3::{exceptions::PyValueError, pyclass, pymethods, PyResult};
 use rand::prelude::Distribution;
-use rand::{thread_rng, RngCore};
+use rand::{rng, RngCore};
 use rand_distr::StandardNormal;
 use smallvec::{SmallVec, ToSmallVec};
 
 use thiserror::Error;
 
+use crate::wrapper::PyTransformAdapt;
+
 type InnerModel = bridgestan::Model<Arc<bridgestan::StanLibrary>>;
 
 #[pyclass]
@@ -53,8 +56,8 @@ impl StanVariable {
     }
 
     #[getter]
-    fn shape<'py>(&self, py: Python<'py>) -> Bound<'py, PyTuple> {
-        PyTuple::new_bound(py, self.0.shape.iter())
+    fn shape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(py, self.0.shape.iter())
     }
 
     #[getter]
@@ -68,6 +71,7 @@ impl StanVariable {
 pub struct StanModel {
     model: Arc<InnerModel>,
     variables: Vec<Parameter>,
+    transform_adapter: Option<PyTransformAdapt>,
 }
 
 /// Return meta information about the constrained parameters of the model
@@ -136,25 +140,41 @@ fn params(
 #[pymethods]
 impl StanModel {
     #[new]
-    pub fn new(lib: StanLibrary, seed: Option<u32>, data: Option<String>) -> anyhow::Result<Self> {
+    #[pyo3(signature = (lib, seed=None, data=None, transform_adapter=None))]
+    pub fn new(
+        lib: StanLibrary,
+        seed: Option<u32>,
+        data: Option<String>,
+        transform_adapter: Option<Py<PyAny>>,
+    ) -> anyhow::Result<Self> {
         let seed = match seed {
             Some(seed) => seed,
-            None => thread_rng().next_u32(),
+            None => rng().next_u32(),
         };
         let data: Option<CString> = data.map(CString::new).transpose()?;
         let model = Arc::new(
             bridgestan::Model::new(lib.0, data.as_ref(), seed).map_err(anyhow::Error::new)?,
         );
         let variables = params(&model, true, true)?;
-        Ok(StanModel { model, variables })
+        let transform_adapter = transform_adapter.map(PyTransformAdapt::new);
+        Ok(StanModel {
+            model,
+            variables,
+            transform_adapter,
+        })
     }
 
     pub fn variables<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
-        let out = PyDict::new_bound(py);
+        let out = PyDict::new(py);
         let results: Result<Vec<_>, _> = self
             .variables
             .iter()
-            .map(|var| out.set_item(var.name.clone(), StanVariable(var.clone()).into_py(py)))
+            .map(|var| {
+                out.set_item(
+                    var.name.clone(),
+                    StanVariable(var.clone()).into_pyobject(py)?,
+                )
+            })
             .collect();
         results?;
         Ok(out)
@@ -195,7 +215,10 @@ impl StanModel {
     */
 }
 
-pub struct StanDensity<'model>(&'model InnerModel);
+pub struct StanDensity<'model> {
+    inner: &'model InnerModel,
+    transform_adapter: Option<PyTransformAdapt>,
+}
 
 #[derive(Debug, Error)]
 pub enum StanLogpError {
@@ -203,6 +226,10 @@ pub enum StanLogpError {
     BridgeStan(#[from] bridgestan::BridgeStanError),
     #[error("Bad logp value: {0}")]
     BadLogp(f64),
+    #[error("Python exception: {0}")]
+    PyErr(#[from] PyErr),
+    #[error("Unspecified Error: {0}")]
+    Anyhow(#[from] anyhow::Error),
 }
 
 impl LogpError for StanLogpError {
@@ -213,9 +240,12 @@ impl LogpError for StanLogpError {
 
 impl<'model> CpuLogpFunc for StanDensity<'model> {
     type LogpError = StanLogpError;
+    type TransformParams = Py<PyAny>;
 
     fn logp(&mut self, position: &[f64], grad: &mut [f64]) -> Result<f64, Self::LogpError> {
-        let logp = self.0.log_density_gradient(position, true, true, grad)?;
+        let logp = self
+            .inner
+            .log_density_gradient(position, true, true, grad)?;
         if !logp.is_finite() {
             return Err(StanLogpError::BadLogp(logp));
         }
@@ -223,7 +253,143 @@ impl<'model> CpuLogpFunc for StanDensity<'model> {
     }
 
     fn dim(&self) -> usize {
-        self.0.param_unc_num()
+        self.inner.param_unc_num()
+    }
+
+    fn inv_transform_normalize(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<f64, Self::LogpError> {
+        let logdet = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .inv_transform_normalize(
+                params,
+                untransformed_position,
+                untransformed_gradient,
+                transformed_position,
+                transformed_gradient,
+            )
+            .context("failed inv_transform_normalize")?;
+        Ok(logdet)
+    }
+
+    fn init_from_transformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &mut [f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &[f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<(f64, f64), Self::LogpError> {
+        let adapter = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?;
+
+        let part1 = adapter
+            .init_from_transformed_position_part1(
+                params,
+                untransformed_position,
+                transformed_position,
+            )
+            .context("Failed init_from_transformed_position_part1")?;
+
+        let logp = self.logp(untransformed_position, untransformed_gradient)?;
+
+        let adapter = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?;
+
+        let logdet = adapter
+            .init_from_transformed_position_part2(
+                params,
+                part1,
+                untransformed_gradient,
+                transformed_gradient,
+            )
+            .context("Failed init_from_transformed_position_part2")?;
+        Ok((logp, logdet))
+    }
+
+    fn init_from_untransformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> std::result::Result<(f64, f64), Self::LogpError> {
+        let logp = self
+            .logp(untransformed_position, untransformed_gradient)
+            .context("Failed to call stan logp function")?;
+
+        let logdet = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .inv_transform_normalize(
+                params,
+                untransformed_position,
+                untransformed_gradient,
+                transformed_position,
+                transformed_gradient,
+            )
+            .context("Failed inv_transform_normalize in stan init_from_untransformed_position")?;
+        Ok((logp, logdet))
+    }
+
+    fn update_transformation<'a, R: rand::Rng + ?Sized>(
+        &'a mut self,
+        rng: &mut R,
+        untransformed_positions: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_gradients: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_logp: impl ExactSizeIterator<Item = &'a f64>,
+        params: &'a mut Py<PyAny>,
+    ) -> std::result::Result<(), Self::LogpError> {
+        self.transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .update_transformation(
+                rng,
+                untransformed_positions,
+                untransformed_gradients,
+                untransformed_logp,
+                params,
+            )
+            .context("Failed to update the transformation")?;
+        Ok(())
+    }
+
+    fn new_transformation<R: rand::Rng + ?Sized>(
+        &mut self,
+        rng: &mut R,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        chain: u64,
+    ) -> std::result::Result<Py<PyAny>, Self::LogpError> {
+        let trafo = self
+            .transform_adapter
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .new_transformation(rng, untransformed_position, untransformed_gradient, chain)
+            .context("Could not create transformation adapter")?;
+        Ok(trafo)
+    }
+
+    fn transformation_id(&self, params: &Py<PyAny>) -> std::result::Result<i64, Self::LogpError> {
+        let id = self
+            .transform_adapter
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err("No transformation adapter specified"))?
+            .transformation_id(params)?;
+        Ok(id)
     }
 }
 
@@ -387,7 +553,10 @@ impl Model for StanModel {
     }
 
     fn math(&self) -> anyhow::Result<Self::Math<'_>> {
-        Ok(CpuMath::new(StanDensity(&self.model)))
+        Ok(CpuMath::new(StanDensity {
+            inner: &self.model,
+            transform_adapter: self.transform_adapter.as_ref().map(|v| v.clone()),
+        }))
     }
 
     fn init_position<R: rand::Rng + ?Sized>(
diff --git a/src/wrapper.rs b/src/wrapper.rs
index 6f9ad49..bb1a523 100644
--- a/src/wrapper.rs
+++ b/src/wrapper.rs
@@ -1,6 +1,7 @@
 use std::{
     fmt::Debug,
-    sync::Arc,
+    ops::{Deref, DerefMut},
+    sync::{Arc, Mutex},
     time::{Duration, Instant},
 };
 
@@ -13,17 +14,19 @@ use crate::{
 
 use anyhow::{bail, Context, Result};
 use arrow::array::Array;
+use numpy::{PyArray1, PyReadonlyArray1};
 use nuts_rs::{
-    AdaptOptions, ChainProgress, DiagAdaptExpSettings, DiagGradNutsSettings, LowRankNutsSettings,
-    LowRankSettings, NutsSettings, ProgressCallback, Sampler, SamplerWaitResult, Trace,
+    ChainProgress, DiagGradNutsSettings, LowRankNutsSettings, ProgressCallback, Sampler,
+    SamplerWaitResult, Trace, TransformedNutsSettings,
 };
 use pyo3::{
     exceptions::PyTimeoutError,
     ffi::Py_uintptr_t,
+    intern,
     prelude::*,
     types::{PyList, PyTuple},
 };
-use rand::{thread_rng, RngCore};
+use rand::{rng, RngCore};
 
 #[pyclass]
 struct PyChainProgress(ChainProgress);
@@ -66,100 +69,23 @@ impl PyChainProgress {
     }
 }
 
-#[derive(Clone)]
-enum InnerSettings {
-    LowRank(LowRankSettings),
-    Diag(DiagAdaptExpSettings),
-}
-
 #[pyclass]
 #[derive(Clone)]
 pub struct PyNutsSettings {
-    settings: NutsSettings<()>,
-    adapt: InnerSettings,
+    inner: Settings,
 }
 
+#[derive(Clone, Debug)]
 enum Settings {
     Diag(DiagGradNutsSettings),
     LowRank(LowRankNutsSettings),
-}
-
-// Would be much nicer with
-// https://doc.rust-lang.org/nightly/unstable-book/language-features/type-changing-struct-update.html
-fn combine_settings<T: Debug + Copy + Default>(
-    inner: T,
-    settings: NutsSettings<()>,
-) -> NutsSettings<T> {
-    let adapt = AdaptOptions {
-        dual_average_options: settings.adapt_options.dual_average_options,
-        mass_matrix_options: inner,
-        early_window: settings.adapt_options.early_window,
-        step_size_window: settings.adapt_options.step_size_window,
-        mass_matrix_switch_freq: settings.adapt_options.mass_matrix_switch_freq,
-        early_mass_matrix_switch_freq: settings.adapt_options.early_mass_matrix_switch_freq,
-        mass_matrix_update_freq: settings.adapt_options.mass_matrix_update_freq,
-    };
-    NutsSettings {
-        num_tune: settings.num_tune,
-        num_draws: settings.num_draws,
-        maxdepth: settings.maxdepth,
-        store_gradient: settings.store_gradient,
-        store_unconstrained: settings.store_unconstrained,
-        max_energy_error: settings.max_energy_error,
-        store_divergences: settings.store_divergences,
-        adapt_options: adapt,
-        check_turning: settings.check_turning,
-        num_chains: settings.num_chains,
-        seed: settings.seed,
-    }
-}
-
-fn split_settings<T: Debug + Copy + Default>(settings: NutsSettings<T>) -> (NutsSettings<()>, T) {
-    let adapt_settings = settings.adapt_options;
-    let mass_matrix_settings = adapt_settings.mass_matrix_options;
-
-    let remaining: AdaptOptions<()> = AdaptOptions {
-        dual_average_options: adapt_settings.dual_average_options,
-        mass_matrix_options: (),
-        early_window: adapt_settings.early_window,
-        step_size_window: adapt_settings.step_size_window,
-        mass_matrix_switch_freq: adapt_settings.mass_matrix_switch_freq,
-        early_mass_matrix_switch_freq: adapt_settings.early_mass_matrix_switch_freq,
-        mass_matrix_update_freq: adapt_settings.mass_matrix_update_freq,
-    };
-
-    let settings = NutsSettings {
-        adapt_options: remaining,
-        num_tune: settings.num_tune,
-        num_draws: settings.num_draws,
-        maxdepth: settings.maxdepth,
-        store_gradient: settings.store_gradient,
-        store_unconstrained: settings.store_unconstrained,
-        max_energy_error: settings.max_energy_error,
-        store_divergences: settings.store_divergences,
-        check_turning: settings.check_turning,
-        num_chains: settings.num_chains,
-        seed: settings.seed,
-    };
-
-    (settings, mass_matrix_settings)
+    Transforming(TransformedNutsSettings),
 }
 
 impl PyNutsSettings {
-    fn into_settings(self) -> Settings {
-        match self.adapt {
-            InnerSettings::LowRank(mass_matrix) => {
-                Settings::LowRank(combine_settings(mass_matrix, self.settings))
-            }
-            InnerSettings::Diag(mass_matrix) => {
-                Settings::Diag(combine_settings(mass_matrix, self.settings))
-            }
-        }
-    }
-
     fn new_diag(seed: Option<u64>) -> Self {
         let seed = seed.unwrap_or_else(|| {
-            let mut rng = thread_rng();
+            let mut rng = rng();
             rng.next_u64()
         });
         let settings = DiagGradNutsSettings {
@@ -167,17 +93,14 @@ impl PyNutsSettings {
             ..Default::default()
         };
 
-        let (settings, inner) = split_settings(settings);
-
         Self {
-            settings,
-            adapt: InnerSettings::Diag(inner),
+            inner: Settings::Diag(settings),
         }
     }
 
     fn new_low_rank(seed: Option<u64>) -> Self {
         let seed = seed.unwrap_or_else(|| {
-            let mut rng = thread_rng();
+            let mut rng = rng();
             rng.next_u64()
         });
         let settings = LowRankNutsSettings {
@@ -185,229 +108,467 @@ impl PyNutsSettings {
             ..Default::default()
         };
 
-        let (settings, inner) = split_settings(settings);
+        Self {
+            inner: Settings::LowRank(settings),
+        }
+    }
+
+    fn new_tranform_adapt(seed: Option<u64>) -> Self {
+        let seed = seed.unwrap_or_else(|| {
+            let mut rng = rng();
+            rng.next_u64()
+        });
+        let settings = TransformedNutsSettings {
+            seed,
+            ..Default::default()
+        };
 
         Self {
-            settings,
-            adapt: InnerSettings::LowRank(inner),
+            inner: Settings::Transforming(settings),
         }
     }
 }
 
+// TODO switch to serde to expose all the options...
 #[pymethods]
 impl PyNutsSettings {
     #[staticmethod]
     #[allow(non_snake_case)]
+    #[pyo3(signature = (seed=None))]
     fn Diag(seed: Option<u64>) -> Self {
         PyNutsSettings::new_diag(seed)
     }
 
     #[staticmethod]
     #[allow(non_snake_case)]
+    #[pyo3(signature = (seed=None))]
     fn LowRank(seed: Option<u64>) -> Self {
         PyNutsSettings::new_low_rank(seed)
     }
 
+    #[staticmethod]
+    #[allow(non_snake_case)]
+    #[pyo3(signature = (seed=None))]
+    fn Transform(seed: Option<u64>) -> Self {
+        PyNutsSettings::new_tranform_adapt(seed)
+    }
+
     #[getter]
     fn num_tune(&self) -> u64 {
-        self.settings.num_tune
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_tune,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_tune,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_tune,
+        }
     }
 
     #[setter(num_tune)]
     fn set_num_tune(&mut self, val: u64) {
-        self.settings.num_tune = val
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_tune = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_tune = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_tune = val,
+        }
     }
 
     #[getter]
     fn num_chains(&self) -> usize {
-        self.settings.num_chains
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_chains,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_chains,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_chains,
+        }
     }
 
     #[setter(num_chains)]
     fn set_num_chains(&mut self, val: usize) {
-        self.settings.num_chains = val;
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_chains = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_chains = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_chains = val,
+        }
     }
 
     #[getter]
     fn num_draws(&self) -> u64 {
-        self.settings.num_draws
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_draws,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_draws,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_draws,
+        }
     }
 
     #[setter(num_draws)]
     fn set_num_draws(&mut self, val: u64) {
-        self.settings.num_draws = val;
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.num_draws = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.num_draws = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.num_draws = val,
+        }
     }
 
     #[getter]
-    fn window_switch_freq(&self) -> u64 {
-        self.settings.adapt_options.mass_matrix_switch_freq
+    fn window_switch_freq(&self) -> Result<u64> {
+        match &self.inner {
+            Settings::Diag(nuts_settings) => {
+                Ok(nuts_settings.adapt_options.mass_matrix_switch_freq)
+            }
+            Settings::LowRank(nuts_settings) => {
+                Ok(nuts_settings.adapt_options.mass_matrix_switch_freq)
+            }
+            Settings::Transforming(nuts_settings) => {
+                Ok(nuts_settings.adapt_options.transform_update_freq)
+            }
+        }
     }
 
     #[setter(window_switch_freq)]
-    fn set_window_switch_freq(&mut self, val: u64) {
-        self.settings.adapt_options.mass_matrix_switch_freq = val;
+    fn set_window_switch_freq(&mut self, val: u64) -> Result<()> {
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings.adapt_options.mass_matrix_switch_freq = val;
+                Ok(())
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings.adapt_options.mass_matrix_switch_freq = val;
+                Ok(())
+            }
+            Settings::Transforming(nuts_settings) => {
+                nuts_settings.adapt_options.transform_update_freq = val;
+                Ok(())
+            }
+        }
     }
 
     #[getter]
-    fn early_window_switch_freq(&self) -> u64 {
-        self.settings.adapt_options.early_mass_matrix_switch_freq
+    fn early_window_switch_freq(&self) -> Result<u64> {
+        match &self.inner {
+            Settings::Diag(nuts_settings) => {
+                Ok(nuts_settings.adapt_options.early_mass_matrix_switch_freq)
+            }
+            Settings::LowRank(nuts_settings) => {
+                Ok(nuts_settings.adapt_options.early_mass_matrix_switch_freq)
+            }
+            Settings::Transforming(_) => {
+                bail!("Option early_window_switch_freq not availbale for transformation adaptation")
+            }
+        }
     }
 
     #[setter(early_window_switch_freq)]
-    fn set_early_window_switch_freq(&mut self, val: u64) {
-        self.settings.adapt_options.early_mass_matrix_switch_freq = val;
+    fn set_early_window_switch_freq(&mut self, val: u64) -> Result<()> {
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings.adapt_options.early_mass_matrix_switch_freq = val;
+                Ok(())
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings.adapt_options.early_mass_matrix_switch_freq = val;
+                Ok(())
+            }
+            Settings::Transforming(_) => {
+                bail!("Option early_window_switch_freq not availbale for transformation adaptation")
+            }
+        }
     }
+
     #[getter]
     fn initial_step(&self) -> f64 {
-        self.settings
-            .adapt_options
-            .dual_average_options
-            .initial_step
+        match &self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step
+            }
+            Settings::Transforming(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step
+            }
+        }
     }
 
     #[setter(initial_step)]
     fn set_initial_step(&mut self, val: f64) {
-        self.settings
-            .adapt_options
-            .dual_average_options
-            .initial_step = val
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step = val;
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step = val;
+            }
+            Settings::Transforming(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .initial_step = val;
+            }
+        }
     }
 
     #[getter]
     fn maxdepth(&self) -> u64 {
-        self.settings.maxdepth
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.maxdepth,
+            Settings::LowRank(nuts_settings) => nuts_settings.maxdepth,
+            Settings::Transforming(nuts_settings) => nuts_settings.maxdepth,
+        }
     }
 
     #[setter(maxdepth)]
     fn set_maxdepth(&mut self, val: u64) {
-        self.settings.maxdepth = val
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.maxdepth = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.maxdepth = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.maxdepth = val,
+        }
     }
 
     #[getter]
     fn store_gradient(&self) -> bool {
-        self.settings.store_gradient
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_gradient,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_gradient,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_gradient,
+        }
     }
 
     #[setter(store_gradient)]
     fn set_store_gradient(&mut self, val: bool) {
-        self.settings.store_gradient = val;
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_gradient = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_gradient = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_gradient = val,
+        }
     }
 
     #[getter]
     fn store_unconstrained(&self) -> bool {
-        self.settings.store_unconstrained
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_unconstrained,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_unconstrained,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_unconstrained,
+        }
     }
 
     #[setter(store_unconstrained)]
     fn set_store_unconstrained(&mut self, val: bool) {
-        self.settings.store_unconstrained = val;
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_unconstrained = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_unconstrained = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_unconstrained = val,
+        }
     }
 
     #[getter]
     fn store_divergences(&self) -> bool {
-        self.settings.store_divergences
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_divergences,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_divergences,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_divergences,
+        }
     }
 
     #[setter(store_divergences)]
     fn set_store_divergences(&mut self, val: bool) {
-        self.settings.store_divergences = val;
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.store_divergences = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.store_divergences = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.store_divergences = val,
+        }
     }
 
     #[getter]
     fn max_energy_error(&self) -> f64 {
-        self.settings.max_energy_error
+        match &self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.max_energy_error,
+            Settings::LowRank(nuts_settings) => nuts_settings.max_energy_error,
+            Settings::Transforming(nuts_settings) => nuts_settings.max_energy_error,
+        }
     }
 
     #[setter(max_energy_error)]
     fn set_max_energy_error(&mut self, val: f64) {
-        self.settings.max_energy_error = val
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => nuts_settings.max_energy_error = val,
+            Settings::LowRank(nuts_settings) => nuts_settings.max_energy_error = val,
+            Settings::Transforming(nuts_settings) => nuts_settings.max_energy_error = val,
+        }
     }
 
-    #[setter(target_accept)]
-    fn set_target_accept(&mut self, val: f64) {
-        self.settings
-            .adapt_options
-            .dual_average_options
-            .target_accept = val;
+    #[getter]
+    fn set_target_accept(&self) -> f64 {
+        match &self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept
+            }
+            Settings::Transforming(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept
+            }
+        }
     }
 
-    #[getter]
-    fn target_accept(&self) -> f64 {
-        self.settings
-            .adapt_options
-            .dual_average_options
-            .target_accept
+    #[setter(target_accept)]
+    fn target_accept(&mut self, val: f64) {
+        match &mut self.inner {
+            Settings::Diag(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept = val
+            }
+            Settings::LowRank(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept = val
+            }
+            Settings::Transforming(nuts_settings) => {
+                nuts_settings
+                    .adapt_options
+                    .dual_average_options
+                    .target_accept = val
+            }
+        }
     }
 
     #[getter]
-    fn store_mass_matrix(&self) -> bool {
-        match &self.adapt {
-            InnerSettings::LowRank(low_rank) => low_rank.store_mass_matrix,
-            InnerSettings::Diag(diag) => diag.store_mass_matrix,
+    fn store_mass_matrix(&self) -> Result<bool> {
+        match &self.inner {
+            Settings::LowRank(settings) => {
+                Ok(settings.adapt_options.mass_matrix_options.store_mass_matrix)
+            }
+            Settings::Diag(settings) => {
+                Ok(settings.adapt_options.mass_matrix_options.store_mass_matrix)
+            }
+            Settings::Transforming(_) => {
+                bail!("Option store_mass_matrix not availbale for transformation adaptation")
+            }
         }
     }
 
     #[setter(store_mass_matrix)]
-    fn set_store_mass_matrix(&mut self, val: bool) {
-        match &mut self.adapt {
-            InnerSettings::LowRank(low_rank) => {
-                low_rank.store_mass_matrix = val;
+    fn set_store_mass_matrix(&mut self, val: bool) -> Result<()> {
+        match &mut self.inner {
+            Settings::LowRank(settings) => {
+                settings.adapt_options.mass_matrix_options.store_mass_matrix = val;
+                Ok(())
+            }
+            Settings::Diag(settings) => {
+                settings.adapt_options.mass_matrix_options.store_mass_matrix = val;
+                Ok(())
             }
-            InnerSettings::Diag(diag) => {
-                diag.store_mass_matrix = val;
+            Settings::Transforming(_) => {
+                bail!("Option store_mass_matrix not availbale for transformation adaptation")
             }
         }
     }
 
     #[getter]
     fn use_grad_based_mass_matrix(&self) -> Result<bool> {
-        match &self.adapt {
-            InnerSettings::LowRank(_) => {
-                bail!("grad based mass matrix not available for low-rank adaptation")
+        match &self.inner {
+            Settings::LowRank(_) => {
+                bail!("non-grad based mass matrix not available for low-rank adaptation")
+            }
+            Settings::Transforming(_) => {
+                bail!("non-grad based mass matrix not available for transforming adaptation")
             }
-            InnerSettings::Diag(diag) => Ok(diag.use_grad_based_estimate),
+            Settings::Diag(diag) => Ok(diag
+                .adapt_options
+                .mass_matrix_options
+                .use_grad_based_estimate),
         }
     }
 
     #[setter(use_grad_based_mass_matrix)]
     fn set_use_grad_based_mass_matrix(&mut self, val: bool) -> Result<()> {
-        match &mut self.adapt {
-            InnerSettings::LowRank(_) => {
-                bail!("grad based mass matrix not available for low-rank adaptation")
+        match &mut self.inner {
+            Settings::LowRank(_) => {
+                bail!("non-grad based mass matrix not available for low-rank adaptation");
+            }
+            Settings::Transforming(_) => {
+                bail!("non-grad based mass matrix not available for transforming adaptation");
             }
-            InnerSettings::Diag(diag) => {
-                diag.use_grad_based_estimate = val;
+            Settings::Diag(diag) => {
+                diag.adapt_options
+                    .mass_matrix_options
+                    .use_grad_based_estimate = val;
             }
         }
         Ok(())
     }
 
     #[getter]
-    fn mass_matrix_switch_freq(&self) -> u64 {
-        self.settings.adapt_options.mass_matrix_switch_freq
+    fn mass_matrix_switch_freq(&self) -> Result<u64> {
+        match &self.inner {
+            Settings::Diag(settings) => Ok(settings.adapt_options.mass_matrix_switch_freq),
+            Settings::LowRank(settings) => Ok(settings.adapt_options.mass_matrix_switch_freq),
+            Settings::Transforming(_) => {
+                bail!("mass_matrix_switch_freq not available for transforming adaptation");
+            }
+        }
     }
 
     #[setter(mass_matrix_switch_freq)]
-    fn set_mass_matrix_switch_freq(&mut self, val: u64) {
-        self.settings.adapt_options.mass_matrix_switch_freq = val;
+    fn set_mass_matrix_switch_freq(&mut self, val: u64) -> Result<()> {
+        match &mut self.inner {
+            Settings::Diag(settings) => settings.adapt_options.mass_matrix_switch_freq = val,
+            Settings::LowRank(settings) => settings.adapt_options.mass_matrix_switch_freq = val,
+            Settings::Transforming(_) => {
+                bail!("mass_matrix_switch_freq not available for transforming adaptation");
+            }
+        }
+        Ok(())
     }
 
     #[getter]
     fn mass_matrix_eigval_cutoff(&self) -> Result<f64> {
-        match &self.adapt {
-            InnerSettings::LowRank(inner) => Ok(inner.eigval_cutoff),
-            InnerSettings::Diag(_) => {
-                bail!("eigenvalue cutoff not available for diag mass matrix adaptation")
+        match &self.inner {
+            Settings::LowRank(inner) => Ok(inner.adapt_options.mass_matrix_options.eigval_cutoff),
+            Settings::Diag(_) => {
+                bail!("eigenvalue cutoff not available for diag mass matrix adaptation");
+            }
+            Settings::Transforming(_) => {
+                bail!("eigenvalue cutoff not available for transfor adaptation");
             }
         }
     }
 
     #[setter(mass_matrix_eigval_cutoff)]
     fn set_mass_matrix_eigval_cutoff(&mut self, val: f64) -> Result<()> {
-        match &mut self.adapt {
-            InnerSettings::LowRank(inner) => inner.eigval_cutoff = val,
-            InnerSettings::Diag(_) => {
-                bail!("eigenvalue cutoff not available for diag mass matrix adaptation")
+        match &mut self.inner {
+            Settings::LowRank(inner) => inner.adapt_options.mass_matrix_options.eigval_cutoff = val,
+            Settings::Diag(_) => {
+                bail!("eigenvalue cutoff not available for diag mass matrix adaptation");
+            }
+            Settings::Transforming(_) => {
+                bail!("eigenvalue cutoff not available for transfor adaptation");
             }
         }
         Ok(())
@@ -415,21 +576,56 @@ impl PyNutsSettings {
 
     #[getter]
     fn mass_matrix_gamma(&self) -> Result<f64> {
-        match &self.adapt {
-            InnerSettings::LowRank(inner) => Ok(inner.gamma),
-            InnerSettings::Diag(_) => {
-                bail!("gamma not available for diag mass matrix adaptation")
+        match &self.inner {
+            Settings::LowRank(inner) => Ok(inner.adapt_options.mass_matrix_options.gamma),
+            Settings::Diag(_) => {
+                bail!("gamma not available for diag mass matrix adaptation");
+            }
+            Settings::Transforming(_) => {
+                bail!("gamma not available for transform adaptation");
             }
         }
     }
 
     #[setter(mass_matrix_gamma)]
     fn set_mass_matrix_gamma(&mut self, val: f64) -> Result<()> {
-        match &mut self.adapt {
-            InnerSettings::LowRank(inner) => inner.gamma = val,
-            InnerSettings::Diag(_) => {
-                bail!("gamma not available for diag mass matrix adaptation")
+        match &mut self.inner {
+            Settings::LowRank(inner) => {
+                inner.adapt_options.mass_matrix_options.gamma = val;
+            }
+            Settings::Diag(_) => {
+                bail!("gamma not available for diag mass matrix adaptation");
+            }
+            Settings::Transforming(_) => {
+                bail!("gamma not available for transform adaptation");
+            }
+        }
+        Ok(())
+    }
+
+    #[getter]
+    fn train_on_orbit(&self) -> Result<bool> {
+        match &self.inner {
+            Settings::LowRank(_) => {
+                bail!("gamma not available for low rank mass matrix adaptation");
+            }
+            Settings::Diag(_) => {
+                bail!("gamma not available for diag mass matrix adaptation");
             }
+            Settings::Transforming(inner) => Ok(inner.adapt_options.use_orbit_for_training),
+        }
+    }
+
+    #[setter(train_on_orbit)]
+    fn set_train_on_orbit(&mut self, val: bool) -> Result<()> {
+        match &mut self.inner {
+            Settings::LowRank(_) => {
+                bail!("gamma not available for low rank mass matrix adaptation");
+            }
+            Settings::Diag(_) => {
+                bail!("gamma not available for diag mass matrix adaptation");
+            }
+            Settings::Transforming(inner) => inner.adapt_options.use_orbit_for_training = val,
         }
         Ok(())
     }
@@ -442,13 +638,12 @@ pub(crate) enum SamplerState {
 }
 
 #[derive(Clone)]
-#[pyclass]
-pub enum ProgressType {
+enum InnerProgressType {
     Callback {
         rate: Duration,
         n_cores: usize,
         template: String,
-        callback: Py<PyAny>,
+        callback: Arc<Py<PyAny>>,
     },
     Indicatif {
         rate: Duration,
@@ -456,10 +651,14 @@ pub enum ProgressType {
     None {},
 }
 
+#[pyclass]
+#[derive(Clone)]
+pub struct ProgressType(InnerProgressType);
+
 impl ProgressType {
     fn into_callback(self) -> Result<Option<ProgressCallback>> {
-        match self {
-            ProgressType::Callback {
+        match self.0 {
+            InnerProgressType::Callback {
                 callback,
                 rate,
                 n_cores,
@@ -470,11 +669,11 @@ impl ProgressType {
 
                 Ok(Some(callback))
             }
-            ProgressType::Indicatif { rate } => {
+            InnerProgressType::Indicatif { rate } => {
                 let handler = IndicatifHandler::new(rate);
                 Ok(Some(handler.into_callback()?))
             }
-            ProgressType::None {} => Ok(None),
+            InnerProgressType::None {} => Ok(None),
         }
     }
 }
@@ -484,28 +683,28 @@ impl ProgressType {
     #[staticmethod]
     fn indicatif(rate: u64) -> Self {
         let rate = Duration::from_millis(rate);
-        ProgressType::Indicatif { rate }
+        ProgressType(InnerProgressType::Indicatif { rate })
     }
 
     #[staticmethod]
     fn none() -> Self {
-        ProgressType::None {}
+        ProgressType(InnerProgressType::None {})
     }
 
     #[staticmethod]
     fn template_callback(rate: u64, template: String, n_cores: usize, callback: Py<PyAny>) -> Self {
         let rate = Duration::from_millis(rate);
-        ProgressType::Callback {
-            callback,
+        ProgressType(InnerProgressType::Callback {
+            callback: Arc::new(callback),
             template,
             n_cores,
             rate,
-        }
+        })
     }
 }
 
 #[pyclass]
-struct PySampler(SamplerState);
+struct PySampler(Mutex<SamplerState>);
 
 #[pymethods]
 impl PySampler {
@@ -517,14 +716,18 @@ impl PySampler {
         progress_type: ProgressType,
     ) -> PyResult<PySampler> {
         let callback = progress_type.into_callback()?;
-        match settings.into_settings() {
+        match settings.inner {
             Settings::LowRank(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
             Settings::Diag(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
+            }
+            Settings::Transforming(settings) => {
+                let sampler = Sampler::new(model, settings, cores, callback)?;
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
         }
     }
@@ -537,14 +740,18 @@ impl PySampler {
         progress_type: ProgressType,
     ) -> PyResult<PySampler> {
         let callback = progress_type.into_callback()?;
-        match settings.into_settings() {
+        match settings.inner {
             Settings::LowRank(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
             Settings::Diag(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
+            }
+            Settings::Transforming(settings) => {
+                let sampler = Sampler::new(model, settings, cores, callback)?;
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
         }
     }
@@ -557,38 +764,45 @@ impl PySampler {
         progress_type: ProgressType,
     ) -> PyResult<PySampler> {
         let callback = progress_type.into_callback()?;
-        match settings.into_settings() {
+        match settings.inner {
             Settings::LowRank(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
             Settings::Diag(settings) => {
                 let sampler = Sampler::new(model, settings, cores, callback)?;
-                Ok(PySampler(SamplerState::Running(sampler)))
+                Ok(PySampler(SamplerState::Running(sampler).into()))
+            }
+            Settings::Transforming(settings) => {
+                let sampler = Sampler::new(model, settings, cores, callback)?;
+                Ok(PySampler(SamplerState::Running(sampler).into()))
             }
         }
     }
 
     fn is_finished(&mut self, py: Python<'_>) -> PyResult<bool> {
         py.allow_threads(|| {
-            let state = std::mem::replace(&mut self.0, SamplerState::Empty);
+            let guard = &mut self.0.lock().expect("Poisond sampler state mutex");
+            let slot = guard.deref_mut();
+
+            let state = std::mem::replace(slot, SamplerState::Empty);
 
             let SamplerState::Running(sampler) = state else {
-                let _ = std::mem::replace(&mut self.0, state);
+                let _ = std::mem::replace(slot, state);
                 return Ok(true);
             };
 
             match sampler.wait_timeout(Duration::from_millis(1)) {
                 SamplerWaitResult::Trace(trace) => {
-                    let _ = std::mem::replace(&mut self.0, SamplerState::Finished(Some(trace)));
+                    let _ = std::mem::replace(slot, SamplerState::Finished(Some(trace)));
                     Ok(true)
                 }
                 SamplerWaitResult::Timeout(sampler) => {
-                    let _ = std::mem::replace(&mut self.0, SamplerState::Running(sampler));
+                    let _ = std::mem::replace(slot, SamplerState::Running(sampler));
                     Ok(false)
                 }
                 SamplerWaitResult::Err(err, trace) => {
-                    let _ = std::mem::replace(&mut self.0, SamplerState::Finished(trace));
+                    let _ = std::mem::replace(slot, SamplerState::Finished(trace));
                     Err(err.into())
                 }
             }
@@ -597,7 +811,12 @@ impl PySampler {
 
     fn pause(&mut self, py: Python<'_>) -> PyResult<()> {
         py.allow_threads(|| {
-            if let SamplerState::Running(ref mut control) = self.0 {
+            if let SamplerState::Running(ref mut control) = self
+                .0
+                .lock()
+                .expect("Poised sampler state mutex")
+                .deref_mut()
+            {
                 control.pause()?
             }
             Ok(())
@@ -606,24 +825,33 @@ impl PySampler {
 
     fn resume(&mut self, py: Python<'_>) -> PyResult<()> {
         py.allow_threads(|| {
-            if let SamplerState::Running(ref mut control) = self.0 {
+            if let SamplerState::Running(ref mut control) = self
+                .0
+                .lock()
+                .expect("Poisond sampler state mutex")
+                .deref_mut()
+            {
                 control.resume()?
             }
             Ok(())
         })
     }
 
+    #[pyo3(signature = (timeout_seconds=None))]
     fn wait(&mut self, py: Python<'_>, timeout_seconds: Option<f64>) -> PyResult<()> {
         py.allow_threads(|| {
+            let guard = &mut self.0.lock().expect("Poisond sampler state mutex");
+            let slot = guard.deref_mut();
+
             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.0, SamplerState::Empty);
+            let state = std::mem::replace(slot, SamplerState::Empty);
 
             let SamplerState::Running(mut control) = state else {
-                let _ = std::mem::replace(&mut self.0, state);
+                let _ = std::mem::replace(slot, state);
                 return Ok(());
             };
 
@@ -664,32 +892,38 @@ impl PySampler {
                 }
             };
 
-            let _ = std::mem::replace(&mut self.0, final_state);
+            let _ = std::mem::replace(slot, final_state);
             retval
         })
     }
 
     fn abort(&mut self, py: Python<'_>) -> PyResult<()> {
         py.allow_threads(|| {
-            let state = std::mem::replace(&mut self.0, SamplerState::Empty);
+            let guard = &mut self.0.lock().expect("Poisond sampler state mutex");
+            let slot = guard.deref_mut();
+
+            let state = std::mem::replace(slot, SamplerState::Empty);
 
             let SamplerState::Running(control) = state else {
-                let _ = std::mem::replace(&mut self.0, state);
+                let _ = std::mem::replace(slot, state);
                 return Ok(());
             };
 
             let (result, trace) = control.abort();
-            let _ = std::mem::replace(&mut self.0, SamplerState::Finished(trace));
+            let _ = std::mem::replace(slot, SamplerState::Finished(trace));
             result?;
             Ok(())
         })
     }
 
     fn extract_results<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
-        let state = std::mem::replace(&mut self.0, SamplerState::Empty);
+        let guard = &mut self.0.lock().expect("Poisond sampler state mutex");
+        let slot = guard.deref_mut();
+
+        let state = std::mem::replace(slot, SamplerState::Empty);
 
         let SamplerState::Finished(trace) = state else {
-            let _ = std::mem::replace(&mut self.0, state);
+            let _ = std::mem::replace(slot, state);
             return Err(anyhow::anyhow!("Sampler is not finished"))?;
         };
 
@@ -703,7 +937,7 @@ impl PySampler {
     }
 
     fn is_empty(&self) -> bool {
-        match self.0 {
+        match self.0.lock().expect("Poisoned sampler state lock").deref() {
             SamplerState::Running(_) => false,
             SamplerState::Finished(_) => false,
             SamplerState::Empty => true,
@@ -712,7 +946,8 @@ impl PySampler {
 
     fn inspect<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
         let trace = py.allow_threads(|| {
-            let SamplerState::Running(ref mut sampler) = self.0 else {
+            let mut guard = self.0.lock().unwrap();
+            let SamplerState::Running(ref mut sampler) = guard.deref_mut() else {
                 return Err(anyhow::anyhow!("Sampler is not running"))?;
             };
 
@@ -723,28 +958,28 @@ impl PySampler {
 }
 
 fn trace_to_list(trace: Trace, py: Python<'_>) -> PyResult<Bound<'_, PyList>> {
-    let list = PyList::new_bound(
+    let list = PyList::new(
         py,
         trace
             .chains
             .into_iter()
             .map(|chain| {
-                Ok(PyTuple::new_bound(
+                Ok(PyTuple::new(
                     py,
                     [
                         export_array(py, chain.draws)?,
                         export_array(py, chain.stats)?,
                     ]
                     .into_iter(),
-                ))
+                )?)
             })
             .collect::<Result<Vec<_>>>()?,
-    );
+    )?;
     Ok(list)
 }
 
 fn export_array(py: Python<'_>, data: Arc<dyn Array>) -> PyResult<PyObject> {
-    let pa = py.import_bound("pyarrow")?;
+    let pa = py.import("pyarrow")?;
     let array = pa.getattr("Array")?;
 
     let data = data.into_data();
@@ -755,12 +990,252 @@ fn export_array(py: Python<'_>, data: Arc<dyn Array>) -> PyResult<PyObject> {
         .call_method1(
             "_import_from_c",
             (
-                (&data as *const _ as Py_uintptr_t).into_py(py),
-                (&schema as *const _ as Py_uintptr_t).into_py(py),
+                (&data as *const _ as Py_uintptr_t).into_pyobject(py)?,
+                (&schema as *const _ as Py_uintptr_t).into_pyobject(py)?,
             ),
         )
         .context("Could not import arrow trace in python")?;
-    Ok(data.into_py(py))
+    Ok(data.unbind())
+}
+
+#[pyclass]
+#[derive(Debug, Clone)]
+pub struct PyTransformAdapt(Arc<Py<PyAny>>);
+
+#[pymethods]
+impl PyTransformAdapt {
+    #[new]
+    pub fn new(adapter: Py<PyAny>) -> Self {
+        Self(Arc::new(adapter))
+    }
+}
+
+impl PyTransformAdapt {
+    pub fn inv_transform_normalize(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> Result<f64> {
+        Python::with_gil(|py| {
+            let untransformed_position = PyArray1::from_slice(py, untransformed_position);
+            let untransformed_gradient = PyArray1::from_slice(py, untransformed_gradient);
+
+            let output = params
+                .getattr(py, intern!(py, "inv_transform"))
+                .context("Could not access attribute inv_transform")?
+                .call1(py, (untransformed_position, untransformed_gradient))
+                .context("Failed to call adapter.inv_transform")?;
+            let (logdet, transformed_position_out, transformed_gradient_out): (
+                f64,
+                PyReadonlyArray1<f64>,
+                PyReadonlyArray1<f64>,
+            ) = output
+                .extract(py)
+                .context("Execpected results from adapter.inv_transform")?;
+
+            if !transformed_position_out
+                .as_slice()?
+                .iter()
+                .all(|&x| x.is_finite())
+            {
+                bail!("Transformed position is not finite");
+            }
+            if !transformed_gradient_out
+                .as_slice()?
+                .iter()
+                .all(|&x| x.is_finite())
+            {
+                bail!("Transformed position is not finite");
+            }
+
+            transformed_position.copy_from_slice(
+                transformed_position_out
+                    .as_slice()
+                    .context("Could not copy transformed_position")?,
+            );
+
+            transformed_gradient.copy_from_slice(
+                transformed_gradient_out
+                    .as_slice()
+                    .context("Could not copy transformed_gradient")?,
+            );
+            Ok(logdet)
+        })
+    }
+
+    pub fn init_from_transformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &mut [f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &[f64],
+        transformed_gradient: &mut [f64],
+    ) -> Result<(f64, f64)> {
+        Python::with_gil(|py| {
+            let transformed_position = PyArray1::from_slice(py, transformed_position);
+
+            let output = params
+                .getattr(py, intern!(py, "init_from_transformed_position"))?
+                .call1(py, (transformed_position,))?;
+            let (
+                logp,
+                logdet,
+                untransformed_position_out,
+                untransformed_gradient_out,
+                transformed_gradient_out,
+            ): (
+                f64,
+                f64,
+                PyReadonlyArray1<f64>,
+                PyReadonlyArray1<f64>,
+                PyReadonlyArray1<f64>,
+            ) = output.extract(py)?;
+
+            untransformed_position.copy_from_slice(untransformed_position_out.as_slice()?);
+            untransformed_gradient.copy_from_slice(untransformed_gradient_out.as_slice()?);
+            transformed_gradient.copy_from_slice(transformed_gradient_out.as_slice()?);
+            Ok((logp, logdet))
+        })
+    }
+
+    pub fn init_from_transformed_position_part1(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &mut [f64],
+        transformed_position: &[f64],
+    ) -> Result<Py<PyAny>> {
+        Python::with_gil(|py| {
+            let transformed_position = PyArray1::from_slice(py, transformed_position);
+
+            let output = params
+                .getattr(py, intern!(py, "init_from_transformed_position_part1"))?
+                .call1(py, (transformed_position,))?;
+            let (untransformed_position_out, part1): (PyReadonlyArray1<f64>, Py<PyAny>) =
+                output.extract(py)?;
+
+            untransformed_position.copy_from_slice(untransformed_position_out.as_slice()?);
+            Ok(part1)
+        })
+    }
+
+    pub fn init_from_transformed_position_part2(
+        &mut self,
+        params: &Py<PyAny>,
+        part1: Py<PyAny>,
+        untransformed_gradient: &[f64],
+        transformed_gradient: &mut [f64],
+    ) -> Result<f64> {
+        Python::with_gil(|py| {
+            let untransformed_gradient = PyArray1::from_slice(py, untransformed_gradient);
+
+            let output = params
+                .getattr(py, intern!(py, "init_from_transformed_position_part2"))?
+                .call1(py, (part1, untransformed_gradient))?;
+            let (logdet, transformed_gradient_out): (f64, PyReadonlyArray1<f64>) =
+                output.extract(py)?;
+
+            transformed_gradient.copy_from_slice(transformed_gradient_out.as_slice()?);
+            Ok(logdet)
+        })
+    }
+
+    pub fn init_from_untransformed_position(
+        &mut self,
+        params: &Py<PyAny>,
+        untransformed_position: &[f64],
+        untransformed_gradient: &mut [f64],
+        transformed_position: &mut [f64],
+        transformed_gradient: &mut [f64],
+    ) -> Result<(f64, f64)> {
+        Python::with_gil(|py| {
+            let untransformed_position = PyArray1::from_slice(py, untransformed_position);
+
+            let output = params
+                .getattr(py, intern!(py, "init_from_untransformed_position"))
+                .context("No attribute init_from_untransformed_position")?
+                .call1(py, (untransformed_position,))
+                .context("Failed adapter.init_from_untransformed_position")?;
+            let (
+                logp,
+                logdet,
+                untransformed_gradient_out,
+                transformed_position_out,
+                transformed_gradient_out,
+            ): (
+                f64,
+                f64,
+                PyReadonlyArray1<f64>,
+                PyReadonlyArray1<f64>,
+                PyReadonlyArray1<f64>,
+            ) = output
+                .extract(py)
+                .context("Unexpected return value of init_from_untransformed_position")?;
+
+            untransformed_gradient.copy_from_slice(untransformed_gradient_out.as_slice()?);
+            transformed_position.copy_from_slice(transformed_position_out.as_slice()?);
+            transformed_gradient.copy_from_slice(transformed_gradient_out.as_slice()?);
+            Ok((logp, logdet))
+        })
+    }
+
+    pub fn update_transformation<'a, R: rand::Rng + ?Sized>(
+        &'a mut self,
+        rng: &mut R,
+        untransformed_positions: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_gradients: impl ExactSizeIterator<Item = &'a [f64]>,
+        untransformed_logp: impl ExactSizeIterator<Item = &'a f64>,
+        params: &'a mut Py<PyAny>,
+    ) -> Result<()> {
+        Python::with_gil(|py| {
+            let positions = PyList::new(
+                py,
+                untransformed_positions.map(|pos| PyArray1::from_slice(py, pos)),
+            )?;
+            let gradients = PyList::new(
+                py,
+                untransformed_gradients.map(|grad| PyArray1::from_slice(py, grad)),
+            )?;
+
+            let logps = PyArray1::from_iter(py, untransformed_logp.copied());
+            let seed = rng.next_u64();
+
+            params
+                .getattr(py, intern!(py, "update"))?
+                .call1(py, (seed, positions, gradients, logps))?;
+            Ok(())
+        })
+    }
+
+    pub fn new_transformation<R: rand::Rng + ?Sized>(
+        &mut self,
+        rng: &mut R,
+        untransformed_position: &[f64],
+        untransformed_gradient: &[f64],
+        chain: u64,
+    ) -> Result<Py<PyAny>> {
+        Python::with_gil(|py| {
+            let position = PyArray1::from_slice(py, untransformed_position);
+            let gradient = PyArray1::from_slice(py, untransformed_gradient);
+
+            let seed = rng.next_u64();
+
+            let transformer = self.0.call1(py, (seed, position, gradient, chain))?;
+
+            Ok(transformer)
+        })
+    }
+
+    pub fn transformation_id(&self, params: &Py<PyAny>) -> Result<i64> {
+        Python::with_gil(|py| {
+            let id: i64 = params
+                .getattr(py, intern!(py, "transformation_id"))?
+                .extract(py)?;
+            Ok(id)
+        })
+    }
 }
 
 /// A Python module implemented in Rust.