Scan related error using nutpie for an autoregressive model

Questions
, , ,
1

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):

2

Update

(I can’t edit the previous post).
I’ve tried using updated versions and now I get the same error but even before considering nutpie. Just calling

with ar_model:
    idata = pm.sample()

yields the error in the original post with versions:

  • pymc==5.20.0
  • pytensor==2.26.4
  • python==3.12.8

I thought that perhaps the issue had to do with using a LogNormal inside the ar step calculation, since that’s the main deviation from the examples and tutorials I’ve read, but changing to plain Normal (which should fail but at a different point, because it could lead to a negative probability being thrown to the geometric) does not fix it.

3

This is a cryptic one, so here is an equally cryptic solution.

Specify the shape of rho:

rho = pm.Beta("rho", 5, 1, dims=("lags",), shape=(1,))

Or more robust:

rho = pm.Beta("rho", 5, 1, dims=("lags",), shape=(lags,))

Coordinates with length 1 can be tricky because PyTensor treats unknown dimensions and dimensions of length 1 very differently, and doesn’t like to interchange between them.

When PyMC compiles the logp+dlogp, it swaps from unknown dimensions to the static shapes (since dims are not allowed to change during sampling), and for very in-the-weeds reasons the scan used in the logp does not like this change. Where it used to see an unknown length variable, it now sees a length 1, boo-oh!

If you had started with a model of two lags, this shouldn’t have raised an error.

3 Likes
4

Wow!! Thank you very much! There’s no way I would’ve figured this out on my own. I may come back with some issues I’m still encountering with nutpie (once I’ve experimented a little bit more), but for the time being I’m happy sticking to standard sampling.

(Edit: I don’t know why I don’t see the option of marking your reply as solution, I usually find it next to the “like”, “bookmark”, and so on buttons).

2 Likes
5

Thanks :pray:

The checkmark was because the topic category was not a question. I did it now

1 Like