Nutpie sampling issue

I set up a new environment and installed it with conda/mamba, and ran it in that environment. You can see the packages below (sorry if a lot is redundant I’m just not sure what will be useful to see):

# packages in environment at /root/miniforge3/envs/nutpie_experiment:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                 conda_forge    conda-forge
_openmp_mutex             4.5                  2_kmp_llvm    conda-forge
anyio                     4.3.0              pyhd8ed1ab_0    conda-forge
argon2-cffi               23.1.0             pyhd8ed1ab_0    conda-forge
argon2-cffi-bindings      21.2.0          py311h459d7ec_4    conda-forge
arrow                     1.3.0              pyhd8ed1ab_0    conda-forge
arviz                     0.18.0             pyhd8ed1ab_0    conda-forge
asttokens                 2.4.1              pyhd8ed1ab_0    conda-forge
async-lru                 2.0.4              pyhd8ed1ab_0    conda-forge
atk-1.0                   2.38.0               h04ea711_2    conda-forge
attrs                     23.2.0             pyh71513ae_0    conda-forge
aws-c-auth                0.7.18               he0b1f16_0    conda-forge
aws-c-cal                 0.6.11               heb1d5e4_0    conda-forge
aws-c-common              0.9.15               hd590300_0    conda-forge
aws-c-compression         0.2.18               hce8ee76_3    conda-forge
aws-c-event-stream        0.4.2                h01f5eca_8    conda-forge
aws-c-http                0.8.1               hdb68c23_10    conda-forge
aws-c-io                  0.14.7               hbfbeace_6    conda-forge
aws-c-mqtt                0.10.4               h50844eb_0    conda-forge
aws-c-s3                  0.5.7                h6be9164_2    conda-forge
aws-c-sdkutils            0.1.15               hce8ee76_3    conda-forge
aws-checksums             0.1.18               hce8ee76_3    conda-forge
aws-crt-cpp               0.26.8               h2150271_2    conda-forge
aws-sdk-cpp               1.11.267             hddb5a97_7    conda-forge
babel                     2.14.0             pyhd8ed1ab_0    conda-forge
beautifulsoup4            4.12.3             pyha770c72_0    conda-forge
binutils_impl_linux-64    2.40                 ha885e6a_0    conda-forge
binutils_linux-64         2.40                 hdade7a5_3    conda-forge
blas                      2.122                  openblas    conda-forge
blas-devel                3.9.0           22_linux64_openblas    conda-forge
bleach                    6.1.0              pyhd8ed1ab_0    conda-forge
brotli                    1.1.0                hd590300_1    conda-forge
brotli-bin                1.1.0                hd590300_1    conda-forge
brotli-python             1.1.0           py311hb755f60_1    conda-forge
bzip2                     1.0.8                hd590300_5    conda-forge
c-ares                    1.28.1               hd590300_0    conda-forge
ca-certificates           2024.2.2             hbcca054_0    conda-forge
cached-property           1.5.2                hd8ed1ab_1    conda-forge
cached_property           1.5.2              pyha770c72_1    conda-forge
cachetools                5.3.3              pyhd8ed1ab_0    conda-forge
cairo                     1.18.0               h3faef2a_0    conda-forge
certifi                   2024.2.2           pyhd8ed1ab_0    conda-forge
cffi                      1.16.0          py311hb3a22ac_0    conda-forge
charset-normalizer        3.3.2              pyhd8ed1ab_0    conda-forge
cloudpickle               3.0.0              pyhd8ed1ab_0    conda-forge
comm                      0.2.2              pyhd8ed1ab_0    conda-forge
cons                      0.4.6              pyhd8ed1ab_0    conda-forge
contourpy                 1.2.1           py311h9547e67_0    conda-forge
cycler                    0.12.1             pyhd8ed1ab_0    conda-forge
debugpy                   1.8.1           py311hb755f60_0    conda-forge
decorator                 5.1.1              pyhd8ed1ab_0    conda-forge
defusedxml                0.7.1              pyhd8ed1ab_0    conda-forge
dm-tree                   0.1.8           py311h7b78aeb_4    conda-forge
entrypoints               0.4                pyhd8ed1ab_0    conda-forge
etuples                   0.3.9              pyhd8ed1ab_0    conda-forge
exceptiongroup            1.2.0              pyhd8ed1ab_2    conda-forge
executing                 2.0.1              pyhd8ed1ab_0    conda-forge
expat                     2.6.2                h59595ed_0    conda-forge
fastprogress              1.0.3              pyhd8ed1ab_0    conda-forge
filelock                  3.14.0             pyhd8ed1ab_0    conda-forge
font-ttf-dejavu-sans-mono 2.37                 hab24e00_0    conda-forge
font-ttf-inconsolata      3.000                h77eed37_0    conda-forge
font-ttf-source-code-pro  2.038                h77eed37_0    conda-forge
font-ttf-ubuntu           0.83                 h77eed37_2    conda-forge
fontconfig                2.14.2               h14ed4e7_0    conda-forge
fonts-conda-ecosystem     1                             0    conda-forge
fonts-conda-forge         1                             0    conda-forge
fonttools                 4.51.0          py311h459d7ec_0    conda-forge
fqdn                      1.5.1              pyhd8ed1ab_0    conda-forge
freetype                  2.12.1               h267a509_2    conda-forge
fribidi                   1.0.10               h36c2ea0_0    conda-forge
gcc                       12.3.0               h915e2ae_6    conda-forge
gcc_impl_linux-64         12.3.0               h58ffeeb_6    conda-forge
gcc_linux-64              12.3.0               h6477408_3    conda-forge
gdk-pixbuf                2.42.11              hb9ae30d_0    conda-forge
gflags                    2.2.2             he1b5a44_1004    conda-forge
giflib                    5.2.2                hd590300_0    conda-forge
glog                      0.7.0                hed5481d_0    conda-forge
graphite2                 1.3.13            h59595ed_1003    conda-forge
graphviz                  9.0.0                h78e8752_1    conda-forge
gtk2                      2.24.33              h280cfa0_4    conda-forge
gts                       0.7.6                h977cf35_4    conda-forge
gxx                       12.3.0               h915e2ae_6    conda-forge
gxx_impl_linux-64         12.3.0               h2a574ab_6    conda-forge
gxx_linux-64              12.3.0               h4a1b8e8_3    conda-forge
h11                       0.14.0             pyhd8ed1ab_0    conda-forge
h2                        4.1.0              pyhd8ed1ab_0    conda-forge
h5netcdf                  1.3.0              pyhd8ed1ab_0    conda-forge
h5py                      3.11.0          nompi_py311hebc2b07_100    conda-forge
harfbuzz                  8.4.0                h3d44ed6_0    conda-forge
hdf5                      1.14.3          nompi_h4f84152_101    conda-forge
hpack                     4.0.0              pyh9f0ad1d_0    conda-forge
httpcore                  1.0.5              pyhd8ed1ab_0    conda-forge
httpx                     0.27.0             pyhd8ed1ab_0    conda-forge
hyperframe                6.0.1              pyhd8ed1ab_0    conda-forge
icu                       73.2                 h59595ed_0    conda-forge
idna                      3.7                pyhd8ed1ab_0    conda-forge
importlib-metadata        7.1.0              pyha770c72_0    conda-forge
importlib_metadata        7.1.0                hd8ed1ab_0    conda-forge
importlib_resources       6.4.0              pyhd8ed1ab_0    conda-forge
ipykernel                 6.29.3             pyhd33586a_0    conda-forge
ipython                   8.22.2             pyh707e725_0    conda-forge
ipywidgets                8.1.2              pyhd8ed1ab_0    conda-forge
isoduration               20.11.0            pyhd8ed1ab_0    conda-forge
jedi                      0.19.1             pyhd8ed1ab_0    conda-forge
jinja2                    3.1.3              pyhd8ed1ab_0    conda-forge
joblib                    1.4.0              pyhd8ed1ab_0    conda-forge
json5                     0.9.25             pyhd8ed1ab_0    conda-forge
jsonpointer               2.4             py311h38be061_3    conda-forge
jsonschema                4.22.0             pyhd8ed1ab_0    conda-forge
jsonschema-specifications 2023.12.1          pyhd8ed1ab_0    conda-forge
jsonschema-with-format-nongpl 4.22.0             pyhd8ed1ab_0    conda-forge
jupyter                   1.0.0             pyhd8ed1ab_10    conda-forge
jupyter-lsp               2.2.5              pyhd8ed1ab_0    conda-forge
jupyter_client            8.6.1              pyhd8ed1ab_0    conda-forge
jupyter_console           6.6.3              pyhd8ed1ab_0    conda-forge
jupyter_core              5.7.2           py311h38be061_0    conda-forge
jupyter_events            0.10.0             pyhd8ed1ab_0    conda-forge
jupyter_server            2.14.0             pyhd8ed1ab_0    conda-forge
jupyter_server_terminals  0.5.3              pyhd8ed1ab_0    conda-forge
jupyterlab                4.1.8              pyhd8ed1ab_0    conda-forge
jupyterlab_pygments       0.3.0              pyhd8ed1ab_1    conda-forge
jupyterlab_server         2.27.1             pyhd8ed1ab_0    conda-forge
jupyterlab_widgets        3.0.10             pyhd8ed1ab_0    conda-forge
kernel-headers_linux-64   2.6.32              he073ed8_17    conda-forge
keyutils                  1.6.1                h166bdaf_0    conda-forge
kiwisolver                1.4.5           py311h9547e67_1    conda-forge
krb5                      1.21.2               h659d440_0    conda-forge
lcms2                     2.16                 hb7c19ff_0    conda-forge
ld_impl_linux-64          2.40                 h55db66e_0    conda-forge
lerc                      4.0.0                h27087fc_0    conda-forge
libabseil                 20240116.2      cxx17_h59595ed_0    conda-forge
libaec                    1.1.3                h59595ed_0    conda-forge
libarrow                  15.0.2           hefa796f_6_cpu    conda-forge
libarrow-acero            15.0.2           hbabe93e_6_cpu    conda-forge
libarrow-dataset          15.0.2           hbabe93e_6_cpu    conda-forge
libarrow-flight           15.0.2           hc4f8a93_6_cpu    conda-forge
libarrow-flight-sql       15.0.2           he4f5ca8_6_cpu    conda-forge
libarrow-gandiva          15.0.2           hc1954e9_6_cpu    conda-forge
libarrow-substrait        15.0.2           he4f5ca8_6_cpu    conda-forge
libblas                   3.9.0           22_linux64_openblas    conda-forge
libbrotlicommon           1.1.0                hd590300_1    conda-forge
libbrotlidec              1.1.0                hd590300_1    conda-forge
libbrotlienc              1.1.0                hd590300_1    conda-forge
libcblas                  3.9.0           22_linux64_openblas    conda-forge
libcrc32c                 1.1.2                h9c3ff4c_0    conda-forge
libcurl                   8.7.1                hca28451_0    conda-forge
libdeflate                1.20                 hd590300_0    conda-forge
libedit                   3.1.20191231         he28a2e2_2    conda-forge
libev                     4.33                 hd590300_2    conda-forge
libevent                  2.1.12               hf998b51_1    conda-forge
libexpat                  2.6.2                h59595ed_0    conda-forge
libffi                    3.4.2                h7f98852_5    conda-forge
libgcc-devel_linux-64     12.3.0             h0223996_106    conda-forge
libgcc-ng                 13.2.0               h77fa898_6    conda-forge
libgd                     2.3.3                h119a65a_9    conda-forge
libgfortran-ng            13.2.0               h69a702a_6    conda-forge
libgfortran5              13.2.0               h43f5ff8_6    conda-forge
libglib                   2.80.0               hf2295e7_6    conda-forge
libgomp                   13.2.0               h77fa898_6    conda-forge
libgoogle-cloud           2.23.0               h9be4e54_1    conda-forge
libgoogle-cloud-storage   2.23.0               hc7a4891_1    conda-forge
libgrpc                   1.62.2               h15f2491_0    conda-forge
libhwloc                  2.10.0          default_h2fb2949_1000    conda-forge
libiconv                  1.17                 hd590300_2    conda-forge
libjpeg-turbo             3.0.0                hd590300_1    conda-forge
liblapack                 3.9.0           22_linux64_openblas    conda-forge
liblapacke                3.9.0           22_linux64_openblas    conda-forge
libllvm14                 14.0.6               hcd5def8_4    conda-forge
libllvm16                 16.0.6               hb3ce162_3    conda-forge
libnghttp2                1.58.0               h47da74e_1    conda-forge
libnl                     3.9.0                hd590300_0    conda-forge
libnsl                    2.0.1                hd590300_0    conda-forge
libopenblas               0.3.27          pthreads_h413a1c8_0    conda-forge
libparquet                15.0.2           hacf5a1f_6_cpu    conda-forge
libpng                    1.6.43               h2797004_0    conda-forge
libprotobuf               4.25.3               h08a7969_0    conda-forge
libre2-11                 2023.09.01           h5a48ba9_2    conda-forge
librsvg                   2.58.0               hadf69e7_1    conda-forge
libsanitizer              12.3.0               hb8811af_6    conda-forge
libsodium                 1.0.18               h36c2ea0_1    conda-forge
libsqlite                 3.45.3               h2797004_0    conda-forge
libssh2                   1.11.0               h0841786_0    conda-forge
libstdcxx-devel_linux-64  12.3.0             h0223996_106    conda-forge
libstdcxx-ng              13.2.0               hc0a3c3a_6    conda-forge
libthrift                 0.19.0               hb90f79a_1    conda-forge
libtiff                   4.6.0                h1dd3fc0_3    conda-forge
libutf8proc               2.8.0                h166bdaf_0    conda-forge
libuuid                   2.38.1               h0b41bf4_0    conda-forge
libwebp                   1.4.0                h2c329e2_0    conda-forge
libwebp-base              1.4.0                hd590300_0    conda-forge
libxcb                    1.15                 h0b41bf4_0    conda-forge
libxcrypt                 4.4.36               hd590300_1    conda-forge
libxml2                   2.12.6               h232c23b_2    conda-forge
libzlib                   1.2.13               hd590300_5    conda-forge
llvm-openmp               18.1.3               h4dfa4b3_0    conda-forge
llvmlite                  0.42.0          py311ha6695c7_1    conda-forge
logical-unification       0.4.6              pyhd8ed1ab_0    conda-forge
lz4-c                     1.9.4                hcb278e6_0    conda-forge
markdown-it-py            3.0.0              pyhd8ed1ab_0    conda-forge
markupsafe                2.1.5           py311h459d7ec_0    conda-forge
matplotlib-base           3.8.4           py311h54ef318_0    conda-forge
matplotlib-inline         0.1.7              pyhd8ed1ab_0    conda-forge
mdurl                     0.1.2              pyhd8ed1ab_0    conda-forge
minikanren                1.0.3              pyhd8ed1ab_0    conda-forge
mistune                   3.0.2              pyhd8ed1ab_0    conda-forge
mkl                       2023.2.0         h84fe81f_50496    conda-forge
mkl-service               2.4.1           py311h6c4e854_0    conda-forge
multipledispatch          0.6.0                      py_0    conda-forge
munkres                   1.1.4              pyh9f0ad1d_0    conda-forge
nbclient                  0.10.0             pyhd8ed1ab_0    conda-forge
nbconvert                 7.16.4               hd8ed1ab_0    conda-forge
nbconvert-core            7.16.4             pyhd8ed1ab_0    conda-forge
nbconvert-pandoc          7.16.4               hd8ed1ab_0    conda-forge
nbformat                  5.10.4             pyhd8ed1ab_0    conda-forge
ncurses                   6.4.20240210         h59595ed_0    conda-forge
nest-asyncio              1.6.0              pyhd8ed1ab_0    conda-forge
notebook                  7.1.3              pyhd8ed1ab_0    conda-forge
notebook-shim             0.2.4              pyhd8ed1ab_0    conda-forge
numba                     0.59.1          py311h96b013e_0    conda-forge
numpy                     1.26.4          py311h64a7726_0    conda-forge
nutpie                    0.10.0          py311h7e2346c_0    conda-forge
openblas                  0.3.27          pthreads_h7a3da1a_0    conda-forge
openjpeg                  2.5.2                h488ebb8_0    conda-forge
openssl                   3.3.0                hd590300_0    conda-forge
orc                       2.0.0                h17fec99_1    conda-forge
overrides                 7.7.0              pyhd8ed1ab_0    conda-forge
packaging                 24.0               pyhd8ed1ab_0    conda-forge
pandas                    2.2.2           py311h320fe9a_0    conda-forge
pandoc                    3.1.13               ha770c72_0    conda-forge
pandocfilters             1.5.0              pyhd8ed1ab_0    conda-forge
pango                     1.52.2               ha41ecd1_0    conda-forge
parso                     0.8.4              pyhd8ed1ab_0    conda-forge
patsy                     0.5.6              pyhd8ed1ab_0    conda-forge
pcre2                     10.43                hcad00b1_0    conda-forge
pexpect                   4.9.0              pyhd8ed1ab_0    conda-forge
pickleshare               0.7.5                   py_1003    conda-forge
pillow                    10.3.0          py311h18e6fac_0    conda-forge
pip                       24.0               pyhd8ed1ab_0    conda-forge
pixman                    0.43.2               h59595ed_0    conda-forge
pkgutil-resolve-name      1.3.10             pyhd8ed1ab_1    conda-forge
platformdirs              4.2.1              pyhd8ed1ab_0    conda-forge
prometheus_client         0.20.0             pyhd8ed1ab_0    conda-forge
prompt-toolkit            3.0.42             pyha770c72_0    conda-forge
prompt_toolkit            3.0.42               hd8ed1ab_0    conda-forge
psutil                    5.9.8           py311h459d7ec_0    conda-forge
pthread-stubs             0.4               h36c2ea0_1001    conda-forge
ptyprocess                0.7.0              pyhd3deb0d_0    conda-forge
pure_eval                 0.2.2              pyhd8ed1ab_0    conda-forge
pyarrow                   15.0.2          py311hd5e4297_6_cpu    conda-forge
pycparser                 2.22               pyhd8ed1ab_0    conda-forge
pygments                  2.17.2             pyhd8ed1ab_0    conda-forge
pymc                      5.14.0               hd8ed1ab_0    conda-forge
pymc-base                 5.14.0             pyhd8ed1ab_0    conda-forge
pymc-marketing            0.5.0              pyhd8ed1ab_0    conda-forge
pyparsing                 3.1.2              pyhd8ed1ab_0    conda-forge
pysocks                   1.7.1              pyha2e5f31_6    conda-forge
pytensor                  2.20.0          py311hb755f60_1    conda-forge
pytensor-base             2.20.0          py311h320fe9a_1    conda-forge
python                    3.11.9          hb806964_0_cpython    conda-forge
python-dateutil           2.9.0              pyhd8ed1ab_0    conda-forge
python-fastjsonschema     2.19.1             pyhd8ed1ab_0    conda-forge
python-graphviz           0.20.3             pyh717bed2_0    conda-forge
python-json-logger        2.0.7              pyhd8ed1ab_0    conda-forge
python-tzdata             2024.1             pyhd8ed1ab_0    conda-forge
python_abi                3.11                    4_cp311    conda-forge
pytz                      2024.1             pyhd8ed1ab_0    conda-forge
pyyaml                    6.0.1           py311h459d7ec_1    conda-forge
pyzmq                     26.0.2          py311h08a0b41_0    conda-forge
qtconsole-base            5.5.1              pyha770c72_0    conda-forge
qtpy                      2.4.1              pyhd8ed1ab_0    conda-forge
rdma-core                 51.0                 hd3aeb46_0    conda-forge
re2                       2023.09.01           h7f4b329_2    conda-forge
readline                  8.2                  h8228510_1    conda-forge
referencing               0.35.1             pyhd8ed1ab_0    conda-forge
requests                  2.31.0             pyhd8ed1ab_0    conda-forge
rfc3339-validator         0.1.4              pyhd8ed1ab_0    conda-forge
rfc3986-validator         0.1.1              pyh9f0ad1d_0    conda-forge
rich                      13.7.1             pyhd8ed1ab_0    conda-forge
rpds-py                   0.18.0          py311h46250e7_0    conda-forge
s2n                       1.4.12               h06160fa_0    conda-forge
scikit-learn              1.4.2           py311hc009520_0    conda-forge
scipy                     1.13.0          py311h64a7726_0    conda-forge
seaborn                   0.13.2               hd8ed1ab_2    conda-forge
seaborn-base              0.13.2             pyhd8ed1ab_2    conda-forge
send2trash                1.8.3              pyh0d859eb_0    conda-forge
setuptools                69.5.1             pyhd8ed1ab_0    conda-forge
six                       1.16.0             pyh6c4a22f_0    conda-forge
snappy                    1.2.0                hdb0a2a9_1    conda-forge
sniffio                   1.3.1              pyhd8ed1ab_0    conda-forge
soupsieve                 2.5                pyhd8ed1ab_1    conda-forge
stack_data                0.6.2              pyhd8ed1ab_0    conda-forge
statsmodels               0.14.1          py311h1f0f07a_0    conda-forge
sysroot_linux-64          2.12                he073ed8_17    conda-forge
tbb                       2021.12.0            h00ab1b0_0    conda-forge
terminado                 0.18.1             pyh0d859eb_0    conda-forge
threadpoolctl             3.5.0              pyhc1e730c_0    conda-forge
tinycss2                  1.3.0              pyhd8ed1ab_0    conda-forge
tk                        8.6.13          noxft_h4845f30_101    conda-forge
tomli                     2.0.1              pyhd8ed1ab_0    conda-forge
toolz                     0.12.1             pyhd8ed1ab_0    conda-forge
tornado                   6.4             py311h459d7ec_0    conda-forge
traitlets                 5.14.3             pyhd8ed1ab_0    conda-forge
types-python-dateutil     2.9.0.20240316     pyhd8ed1ab_0    conda-forge
typing-extensions         4.11.0               hd8ed1ab_0    conda-forge
typing_extensions         4.11.0             pyha770c72_0    conda-forge
typing_utils              0.1.0              pyhd8ed1ab_0    conda-forge
tzdata                    2024a                h0c530f3_0    conda-forge
ucx                       1.15.0               ha691c75_8    conda-forge
umap                      0.1.1                    pypi_0    pypi
uri-template              1.3.0              pyhd8ed1ab_0    conda-forge
urllib3                   2.2.1              pyhd8ed1ab_0    conda-forge
wcwidth                   0.2.13             pyhd8ed1ab_0    conda-forge
webcolors                 1.13               pyhd8ed1ab_0    conda-forge
webencodings              0.5.1              pyhd8ed1ab_2    conda-forge
websocket-client          1.8.0              pyhd8ed1ab_0    conda-forge
wheel                     0.43.0             pyhd8ed1ab_1    conda-forge
widgetsnbextension        4.0.10             pyhd8ed1ab_0    conda-forge
xarray                    2024.3.0           pyhd8ed1ab_0    conda-forge
xarray-einstats           0.7.0              pyhd8ed1ab_0    conda-forge
xorg-kbproto              1.0.7             h7f98852_1002    conda-forge
xorg-libice               1.1.1                hd590300_0    conda-forge
xorg-libsm                1.2.4                h7391055_0    conda-forge
xorg-libx11               1.8.9                h8ee46fc_0    conda-forge
xorg-libxau               1.0.11               hd590300_0    conda-forge
xorg-libxdmcp             1.1.3                h7f98852_0    conda-forge
xorg-libxext              1.3.4                h0b41bf4_2    conda-forge
xorg-libxrender           0.9.11               hd590300_0    conda-forge
xorg-renderproto          0.11.1            h7f98852_1002    conda-forge
xorg-xextproto            7.3.0             h0b41bf4_1003    conda-forge
xorg-xproto               7.0.31            h7f98852_1007    conda-forge
xz                        5.2.6                h166bdaf_0    conda-forge
yaml                      0.2.5                h7f98852_2    conda-forge
zeromq                    4.3.5                h75354e8_3    conda-forge
zipp                      3.17.0             pyhd8ed1ab_0    conda-forge
zlib                      1.2.13               hd590300_5    conda-forge
zstd                      1.5.5                hfc55251_0    conda-forge

The more complex model crashes when trying to sample with nutpie with the error

File "../../root/miniforge3/envs/nutpie_experiment/lib/python3.11/site-packages/nutpie/compile_pymc.py", line 499:
    def extract_shared(x, user_data_):
        <source elided>

        return inner(x, *_shared_tuple)
        ^

And the simplified model crashes with the same error above.