Error with custom distribution after using scan

Questions v5
2

Update

After fixing a simpler version, let’s call it v1 (Scan related error using nutpie for an autoregressive model - #3 by ricardoV94), I’ve come back to this one, let’s call it v2. v1 does work with custom distributions, and I think the key difference is that while v1 has a single initial distribution:

    traffic_init = pm.LogNormal("traffic_init", traffic_mu, traffic_sigma, dims=("instances",))

v2 has a list of possible elements and one is selected depending on the weather at that time:

    traffic_init = [
        pm.LogNormal(
            f"traffic_init_{idx}",
            traffic_weather_mu[weather],
            traffic_weather_sigma[weather],
        )
        for idx, weather in enumerate(weather_first)
    ]

So I think the issue must be here but I can’t find it. I’ve tried wrapping the traffic_init list in a tensor variable using pytensor.tensor.stack() but it does not help. I can get draws just fine:

pm.draw(ar_model["delay"])

it’s when I call ar_model.point_logps() (or pm.sample) that fails.

Versions

  • pymc==5.20.0
  • pytensor==2.26.4