Best Way to Handle a for Loop in pymc4

Questions v5
, ,
8

In general you don’t want to do any intermediate compilation of Aesara functions, because Aesara does lots of optimizations and graph re-writes at compile time to make things run faster. If you compile in chunks, you potentially lose out on optimizations between the chunks.

PyMC will compile the entire model, including your scan, and call it internally when it does its MCMC voodoo. So yes, you want to use result directly in your deterministic function. You can look back over the scan tutorial link to get a good sense of exactly what scan gives you back. It’s a stack of scan step results from t=1 to t=T, as a symbolic tensor (note that the value defined in outputs_info, t=0, is NOT included!), with the “batch” or “step” dimension as the first dimension. So if your function returns a scalar, and you have k steps, you will get a k x 1 vector out. You can continue to operate on it like any other symbolic tensor.

Anyway, a consequence of letting PyMC handle it is that the iscalar values you defined will not be included anywhere, so you need to replace these with the actual values you want to be used for your model. If you need these to be variable, you could wrap the model creation in a function, like:

def build_tweet_model(data, k, v, x0 = 0):
    tweet_model = pm.Model()
    
    with tweet_model:
        alpha = pm.Uniform("alpha", lower=-30, upper=700)
        beta = pm.LogNormal("beta", mu=0, sigma=2)
        delta = pm.Beta("delta", alpha=5, beta=1)

        outputs_info = at.as_tensor_variable(np.asarray(x0, dtype='float64'))

        def inner_func(prior_result):
            return delta * prior_result + data.gamma()

        virality, updates = aesara.scan(fn=inner_func,
                                      outputs_info=outputs_info,
                                      n_steps=k)

        shares = alpha*(pm.math.exp(beta*virality) - 1/(virality+1))
        
    return tweet_model

Now if you wanted to fit this model to several different datasets with different values of k and v, it’s quite easy.

Or, if you wanted to estimate any of these values, you would just put prior over them. You could, for example, define outputs_info = pm.Normal('x0', mu=0, sigma=1), and let the model estimate the initial state. You can’t do that for the number of steps, though, because the shape of the scan result needs to be static. You will need to define that up front.

As for the specific error, it seems like np.isnan is being called somewhere else in your code, but it’s not in this snippet, so I can’t say where your error is occurring. Perhaps something to do with data.gamma()? I couldn’t say more without knowing more.