Using aesara.scan to "double loop" a vector over an matrix

Questions
3

I think you’re over-engineering the problem a bit. Sticking to your example numpy code, what you really want is a a function with inputs X, y, M, P, that loops only over the range P to M, and does a specific computation.

A nice trick when using scan is to use aesara.tensor.arange as the sequence argument, and provide the list of objects to be indexed as non-sequences. This will allow you to collapse the two scans you tried into a single scan. I wrote your numpy example in aesara like this:

y_aes = at.dvector('y')
X_aes = at.dmatrix('X')
P_aes = at.iscalar('P')
M_aes = at.iscalar('M')

results, _ = aesara.scan(lambda j, y, X, P: (y * X[:, (j-P):j]).sum(axis=1),
                         sequences=at.arange(P_aes, M_aes),
                         non_sequences=[y_aes, X_aes, P])

Z_aes = at.concatenate([results], axis=0).T

The main scan is exactly the for-loop in your example code. The only line that I think needs a bit of contemplation is Z_aes = at.concatenate([results], axis=0).T. The lambda function in the scan (which I stole from your code) outputs a single column of computation. Here we are indexing X in such a way that we will get vectors out, so the result object is a list of M-P 1d vectors. When you call concatenate on them, you’ll get back a M-P x N matrix, so you need to transpose it back to N x M-P.

Alternatively, you could try to fuss with shapes in the scan code, so that the lambda function returned N x 1 column vectors.

Verify this gives the right answer:

f_Z = aesara.function([y_aes, X_aes, P_aes, M_aes], [Z_aes])
np.allclose(f_Z(y, X, P, M)[0], Z)
>>> Out: True

Finally, if your ultimate goal is to use these computations in a PyMC model, be aware that you should not compile and evaluate the function using aesara.function. I did that just as a check to show it’s giving the right answer.

7 Likes