Nutpie sampling issue

Questions v5
,
1

Hi all,

I’ve built a HLM and have been able to run the model with the numpyro sampler and the nuts sampler. The results of the model look promising and the fits are good. I am trying to speed it up even further though, and came across the nutpie sampler.

I have followed the docs here, but I fail on this line:

compiled_model = nutpie.compile_pymc_model(pymc_model)

with the error message:

TypingError: Failed in nopython mode pipeline (step: inline inlinable functions)
Failed in untyped mode pipeline (step: handles literal_unroll)
Failed in literal_unroll_subpipeline mode pipeline (step: performs partial type inference)
Untyped global name 'tuple_setitem_literal': Cannot determine Numba type of <class 'function'>

File "path/to/nutpie/compile_pymc.py", line 497:
    def extract_shared(x, user_data_):
        <source elided>
            dat = extract_array(user_data["shared"], index)
            _shared_tuple = tuple_setitem_literal(_shared_tuple, index, dat)
            ^

Is this something anyone else has come across? I am wondering if I need to redefine my model in some way, or if it is a bug given that I am able to run it with other samplers? The same error occurs if I call nuts_sampler = "nutpie" inside of pm.sample().

Also, are there plans to be able to pass the var_names argument for nutpie?

Thanks heaps.

2

What versions of pymc/nutpie/numba do you have installed?

3 
pymc 5.14.0
nutpie 0.10.0
numba 0.59.1
4

The latest numba release is 0.59.1, so if you really have 0.10.0 installed that might be an issue :slight_smile:

But assuming that that was a typo, the error looks like a bug either in the pytensor numba backend, nutpie or numba…
Can you share an example model (preferably small) that fails in this way?

5

@aseyboldt I think @Toph is following nutpie readme example

6

@aseyboldt yes sorry that is a typo - I will edit that comment. I am running numba 0.59.1.

@ricardoV94 I have built my own HLM and am now trying to call it with nutpie as per the docs. The model is quite large so I will try to share a small example model with the same failing.

Currently dealing with a kernel that is crashing because of LLVM ERROR: Symbol not found: __powidf2 when trying to sample with nutpie which I don’t think was the original error I was facing. Will try to reproduce and post here.

7

I’m struggling to reproduce the error with a simplified model, but I am crashing the kernel when trying to sample with nutpie on a simplified model. I keep getting the LLVM ERROR: Symbol not found: __powidf2 with the following simplified model:

with pm.Model(coords=coords) as hlm:
    channel_data = pm.Data("channel_data", X_data, dims=("date", "channel", "state", "cohort"))

    target_value = pm.Data(name="target_y", value=y.values, dims=("date", "state", "cohort"))

    ## INTERCEPT ##
    # Prior#
    mu_a = pm.Normal("mu_a", mu=0.0, sigma=10.0)

    # Non-centered random intercepts - state level
    z_a_s = pm.Normal("z_a_s", mu=0, sigma=1, dims=("state"))
    sigma_a_s = pm.Exponential("sigma_a_s", 5)

    # Non-centered random intercepts - cohort level
    z_a_c = pm.Normal("z_a_c", mu=0, sigma=1, dims=("cohort"))
    sigma_a_c = pm.Exponential("sigma_a_c", 5)

    # Combined#
    alpha = pm.Deterministic(
        "alpha",
        mu_a + z_a_s[:, np.newaxis] * sigma_a_s + z_a_c[np.newaxis, :] * sigma_a_c,
        dims=("state", "cohort"),
    )

    ## Slopes ##
    # Prior#
    mu_b = pm.Normal("mu_b", mu=0.0, sigma=10.0, dims=("channel"))

    # Non-centered random slopes - state level
    z_b_s = pm.Normal("z_b_s", mu=0, sigma=1, dims=("channel", "state"))
    sigma_b_s = pm.Exponential("sigma_b_s", 5)

    # Non-centered random intercepts - cohort level
    z_b_c = pm.Normal("z_b_c", mu=0, sigma=1, dims=("channel", "cohort"))
    sigma_b_c = pm.Exponential("sigma_b_c", 5)

    # combined
    beta = pm.Deterministic(
        "beta",
        mu_b[:, np.newaxis, np.newaxis]
        + z_b_s[:, :, np.newaxis] * sigma_b_s
        + z_b_c[:, np.newaxis, :] * sigma_b_c,
        dims=("channel", "state", "cohort"),
    )

    channel_contribution = pm.Deterministic(
        "beta_channel", var=beta * channel_data, dims=("date", "channel", "state", "cohort")
    )

    # Model error
    sigma_y = pm.Exponential("sigma_y", 1, dims=("state", "cohort"))

    # Expected value
    y_hat = pm.Deterministic(
        "y_hat",
        alpha[np.newaxis, :, :] + channel_contribution.sum(axis=1),
        dims=("date", "state", "cohort"),
    )

    # Data likelihood
    y_like = pm.Normal(
        "y_like",
        mu=y_hat,
        sigma=sigma_y[np.newaxis, :, :],
        observed=target_value,
        dims=("date", "state", "cohort"),
    )

The above model samples fine with numpyro and nuts.

The more complex model does not crash the kernel, but raises the original error.

Is it possible that the simplified model is crashing because of poor package installs (on my part) and the more complex model is not even able to get to that stage and crashes because of something else?

The simplified model crashes almost immediately whilst the complex model spins for a few minutes before throwing the error.

8

Can you add the coords?

9

And this sounds related:

10

How are you installing numba/nutpie? Are you using the conda-forge channel?

11 
coords = {
    "date": dates,
    "channel": channels,
    "state": mn_state,
    "cohort": mn_cohort,
}

where

dates = [ 0  1  2 ... 99]
channels = ['TV_val']
mn_state = Index(['Queensland', 'Victoria', 'WA', 'NT'], dtype='object')
mn_cohort = Index(['Males', 'Females', 'Other'], dtype='object')

If it makes any difference, the more complex model is using an adstock and saturation function on the channel data from the pymc_marketing package.

12

I’m running in a devcontainer - I have pip installed both numba and nutpie/

13

Sounds like it could be an insufficient installation issue then. Can you try setting up a conda/mamba environment and installing via conda-forge channel?

If the problem goes away you know ot was due to how you’re installing it.

14

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.

15

This here might be the underlying issue. If this is the case you should be able to work around the issue until we fix it by removing the channel dimension that is 1, or making that dimension longer.
I really don’t know about the LLVM ERROR: Symbol not found, that really sounds like a numba issue.

16

Latest developments seems to be a Numba bug that it’s showing up in your model. We may have a work-around in nutpie/pytensor sooner than numba fixes it

17

Thanks for the help @aseyboldt and @ricardoV94. Fortunately I am still able to sample with numpyro for the moment. I’ll be sure to check this post if a work around becomes available!