Using PyTensor.scan to loop over multidimensional array

Questions v5
1

Hi all,

for a model I am writing I am trying to implement a loop over the last dimension of a multi-dimensional array A_array (three dimensions).
It is a differential evolution, where the next step always depends on the previous one. Since there is no simple analytical version in my case I am trying to do it like this (simplified model):

A_diff = -at.diff(A_array, axis=2)

    A_new = at.zeros(shape=(at.shape(A_array).eval()))
    A_new = at.set_subtensor(A_new [:,:,0], A_array[:,:,0])

    for i in range(int(at.shape(time).eval())-1):
        b_func = b*A_new [:,:,i]*(A_new [:,:,i]+2*c)

        A_new = at.set_subtensor(A_new [:,:,i+1],  A_new [:,:,i] - (A_diff [:,:,i] + b_func)

return A_new

When I run this loop in numpy it takes about 0.06 seconds, but when A_array is a pytensor.tensor it takes ~ 100 times longer. Why is that?

I have tried to use PyTensor.scan, but b and c are pymc random variables, and somehow my PyTensor.scan implementation did not like that.
Any help of how to speed this loop up or properly write it in PyTensor would be highly appreciated!

2

Looping over pytensor variables takes a long time because it’s not doing what you think it’s doing. Rather than doing a recursive computation, it’s iteratively constructing a computational graph, polluting it with new unnecessary variables and operations at every iteration. You should not use loops in pytensor to do recursive computation.

You should be able to use random variables with scan. What did you try?

3

Thank you for the quick reply!
That makes a lot of sense. The scan implementation I used was with an older version of the model (just one dimensional) and I used the scan(fn: lambda (myfunction), sequences= i ).
But then I got the error, that one of my variable was a random variable and scan couldn’t handle that.
Now I don’t understand the syntax of scan good enough to implement the above code snipped in a useful way. For example, how do I tell scan to only loop over one dimension of a PyTensor tensor?

4

Scan always loops over the first dimension of stuff passed to the sequences argument, so you could permute the axes to put the “batch” dimension first. Another strategy would be to give a pt.arange as the sequence argument and the tensors as non_sequences, then index the tensors in the inner function. Here’s an example of the 2nd strategy, here’s an example of the 1st.

1 Like
5

Okay thank you!
I’ll have a look at both versions…

6

@jessegrabowski:
The issue I seem to have with both version is that A_new[i+1] is informed by the previous step A_new[i]. In my example above I used at.set_subtensor(), but how does that work together with scan?

7

You can use set_subtensor inside the scan function, that’s no problem at all. To recursively pass a variable into scan, use the outputs_info argument. Maybe it will look something like this (untested code):

def step(i, A, A_diff , b, c):
    b_func = b*A [:,:,i]*(A [:,:,i]+2*c)
    A = at.set_subtensor(A[:,:,i+1],  A[:,:,i] - (A_diff [:,:,i] + b_func)
    return A

A_new = at.zeros(shape=(at.shape(A_array).eval()))
A_new = at.set_subtensor(A_new [:,:,0], A_array[:,:,0])

result, updates = pytensor.scan(step, 
                                sequences=at.arange(time.shape[0] - 1),
                                outputs_info = A_new,
                                non_sequences=[A_diff, b, c])

EDIT: I called the argument outputs_info just outputs originally. Oops.

1 Like
8

I tried your suggestion, but the result has 4 dimensions (an additional one for i).
Based on your previous suggestions and this example I used the following:

def step(A_diff, A, b, c):
    b_func = k1*A*(A+2*c)
    return A- (A_diff + b_func)

A_sim = A_array[:,:,0]

result, updates = pytensor.scan(fn=step,
                                        outputs_info=A_sim.T,
                                        sequences=A_diff.T,
                                        non_sequences=[b, c])

However, I noticed that result is then missing the first element of A_sim, so I have to create the final array via:

A_new = at.zeros(shape=(at.shape(A_array).eval()))
A_new = at.set_subtensor(A_new [:,:,0], A_array[:,:,0])
A_new = at.set_subtensor(A_new [:,:,1:], result.T)

Maybe there is a smarter way for this last step, but it runs as supposed now and is much faster! It now takes about 0.2 seconds for a simulation.
Thank you very much for your help @jessegrabowski.

1 Like
9

Nice, I’m glad you got it figured out.

You need to add the initial state yourself, as you discovered. You could just pt.concatenate the result with A_sim[:, :, 0] though; no need to create a new array and fill it. Also you don’t need to evaluate at.shape(A_array), it’s fine to give symbolic tensors as shapes.

10

So can I just do:

A_new = pt.concatenate([A_sim, result],axis=2)

Or is the syntax different?

11

Something like that should work. By default I think it would be on the first axis (axis=0) though, because scan goes over the first axis. You’d have to check the shapes.

12

Yes that works, but I need to use:

A_sim = A_sim.dimshuffle('x',1, 0)
A_new = pt.concatenate([A_sim, result], axis=0)

because A_sim is a subtensor of A_array and has less dimensions.

1 Like
13

Hi @jessegrabowski:
one additional question came up after using this for a while now:
Is there a way to have two matrices in outputs_info?
What is the syntax for that, since there pytensor.scan only outputs result, updates ?

14

You can have your inner function return a tuple for results. For example:

A = pt.dmatrix('A')
B = pt.dmatrix('B')
C = pt.dmatrix('C')

def step(A, B, C):
    return A @ C, B @ C

results, updates = pytensor.scan(step, outputs_info=[A, B], non_sequences=[C], n_steps=10)

results will be a list of 2 (10, None, None) symbolic tensors.

1 Like