Scan related error using nutpie for an autoregressive model

Hi all,

I’m getting the following error when trying to use nutpie (“standard” sampling works) for an autoregressive model. Any contribution is appreciated.

ValueError: Argument Sigmoid.0 given to the scan node is not compatible with its corresponding loop function variable *2-<Vector(float64, shape=(?,))>

MRE

import numpy as np
import pandas as pd
import nutpie
import pymc as pm
from pymc.pytensorf import collect_default_updates
import pytensor

# Create fake data
lags = 1
n_stops = 5
steps = n_stops - lags
n_trips = 2
stop_sequence = np.arange(n_stops, dtype=int)
trips = np.arange(n_trips, dtype=int)
delay_delta = np.array(
    [
        [1, 1, 1, 2, 4],
        [2, 2, 1, 1, 1],
    ]
)

delay_df = pd.DataFrame(delay_delta, columns=stop_sequence)

coords = {
    "lags": range(-lags, 0),
    "stops": stop_sequence,
    "steps": stop_sequence[:-lags],
    "instances": trips,
}


def get_ar_dist(ar_init, rho, mu, sigma, n_steps, size):
    def get_next_step(traffic_t1, rho, mu, sigma):
        traffic_ref = pm.LogNormal.dist(mu=mu, sigma=sigma)
        traffic = traffic_t1 * rho[0] + traffic_ref * (1 - rho[0])
        return traffic, collect_default_updates([traffic])
    #
    ar_steps, _ = pytensor.scan(
        fn=get_next_step,
        outputs_info=[{"initial": ar_init, "taps": [-lags]}],
        # outputs_info=[{"initial": ar_init, "taps": range(-lags, 0)}],
        non_sequences=[rho, mu, sigma],
        n_steps=n_steps,
        strict=True,
    )
    return ar_steps


with pm.Model(coords=coords) as ar_model:
    # Data
    delay_t = pm.Data(
        "delay_t_obs",
        delay_df.values[:, 1:],
        dims=("instances", "steps"),
    )
    # Parameters
    rho = pm.Beta("rho", 5, 1, dims=("lags",))  # traffic persistence
    traffic_mu = pm.Normal("traffic_mu", 0, 1)
    traffic_sigma = pm.HalfNormal("traffic_sigma", 1)
    # AR Model
    traffic_0 = pm.LogNormal("traffic_init", traffic_mu, traffic_sigma, dims=("instances",))
    traffic_t = pm.CustomDist(
        f"traffic_t",
        traffic_0,
        rho,
        traffic_mu,
        traffic_sigma,
        steps,
        dist=get_ar_dist,
        dims=(
            f"steps",
            f"instances",
        ),  # scan iterates over leftmost dimension
    )
    # Likelihood
    pm.Geometric(
        "delay",
        p=1 - pm.math.exp(-traffic_t.T),
        observed=delay_t,
        dims=("instances", "steps"),
    )


compiled_model = nutpie.compile_pymc_model(ar_model)
# trace_pymc = nutpie.sample(compiled_model)

Versions

pymc==5.17.0
pytensor==2.25.5
nutpie==0.13.2
python==3.10.14

Related (actually weakly related, there the error can be avoided by avoiding custom distributions but I got the impression that scan had something to do too):

Scan related error using nutpie for an autoregressive model (original) (raw)

MRE

Versions