Generalising a Custom AR(1) Model to Multiple Dimensions

Questions v5
3

Hi there, you’re right that I want multiple ARs rather than a multivariate (vector) AR. I did try and use pm.AR(), but I found that it can’t cope with the case where rho and sigma depend on both the dimension and the timestep. In other words, I don’t think it can cope with rho and sigma being 2D (# dimensions by # timesteps).

This code, for instance gives an error:

import pymc as pm
import pytensor.tensor as pt

x = pm.AR.dist(
  init_dist=pm.DiracDelta.dist(0.0), 
  rho=pt.exp(pm.Normal.dist(shape=(3, 39))), # 40 times means 39 timesteps
  sigma = pt.exp(pm.Normal.dist(shape=(3, 39))),
  ar_order=1,
  shape=(3, 40), # 3 dimensions, 40 samples
)

pm.draw(x)

but if instead I use rho=pt.exp(pm.Normal.dist(shape=(3, 1))), sigma=pt.exp(pm.Normal.dist()),, it runs fine.

As an aside, it complains if I use rho=pt.exp(pm.Normal.dist(shape=(3, 1))), sigma=pt.exp(pm.Normal.dist(shape=(3, 1))),, but from what you’re saying, perhaps it shouldn’t?

Either way, this apparent limitation lead me to try using CustomDist and scan and I’m having the problem in the original post. Any thoughts?