Using theano.tensor.roll in Model

Hi everyone,

Was hoping somebody could shed some light into a simple model I am trying to run. I would like to use a time-series x(t) as a predictor for another time-series y(t). My model would be something like this: y(t) \sim x(t - T) + \varepsilon, T \sim Uniform(a, b), \varepsilon\sim Normal. So I setup a very simple example:

import numpy as np
import pymc3 as pm
import theano.tensor as tt

x = np.cos(np.linspace(0, 2*np.pi)) * np.linspace(0, 2*np.pi)
y = np.roll(x, shift=-5) + 0.2 * np.random.normal(size=len(x))

def get_lag_model(predictor, observed):
    with pm.Model() as lag_model:
        sigma = pm.HalfNormal("sigma", 1, testval=1.0)
        lag = pm.DiscreteUniform("lag", lower = -10, upper=-1)
        y = pm.Normal("y", mu=tt.roll(predictor, shift=lag, axis=0), sigma=sigma, observed=observed)
    return lag_model

lag_model = get_lag_model(x, y)

When I call the get_lag_model function, I get the following error:

~/.local/lib/python3.8/site-packages/theano/tensor/basic.py in roll(x, shift, axis)
   4590     end_slice = slice(0, -shift)
   4591     end_list = [allslice] * axis + [end_slice] + [allslice] * (x.ndim - axis - 1)
-> 4592     return join(axis, x.__getitem__(tuple(front_list)), x.__getitem__(tuple(end_list)))
   4593 
   4594 

TypeError: slice indices must be integers or None or have an __index__ method

Although I superficially understand what is going on, I am not sure how to fix it. My understanding is pymc3 distributions return a class which will “act” like their dtypes (in this case int64) to theano. But I believe in this case, this is not working properly.

I am happy to hear any suggestions on how to solve this problem, including other forms of parametrizing the problem!

Cheers!