Nutpie sampling issue

Questions v5
,
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.