# Using PyTensor.scan to loop over multidimensional array

**URL:** <https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976>\
**Category:** v5\
**Created:** [April 25, 2023, 5:48pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976 "2023-04-25T17:48:39Z")\
**Posts on this page:** 14\
**Page:** 1

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [April 25, 2023, 5:48pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/1 "2023-04-25T17:48:39Z")

</div>

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):

```auto
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!

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [April 25, 2023, 6:01pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/2 "2023-04-25T18:01:02Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [April 25, 2023, 7:41pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/3 "2023-04-25T19:41:52Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [April 25, 2023, 7:49pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/4 "2023-04-25T19:49:33Z")

</div>

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](https://discourse.pymc.io/t/using-aesara-scan-to-double-loop-a-vector-over-an-matrix/10305/3) of the 2nd strategy, [here’s an example](https://discourse.pymc.io/t/issue-passing-multiple-sequences-to-aesara-scan/10709) of the 1st.

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [April 25, 2023, 8:00pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/5 "2023-04-25T20:00:11Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [May 4, 2023, 4:42pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/6 "2023-05-04T16:42:53Z")

</div>

@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?

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [May 4, 2023, 8:00pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/7 "2023-05-04T20:00:55Z")

</div>

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):

```auto
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.

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [May 5, 2023, 8:03am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/8 "2023-05-05T08:03:43Z")

</div>

I tried your suggestion, but the `result` has 4 dimensions (an additional one for i).  
Based on your previous suggestions and [this example](https://pytensor.readthedocs.io/en/latest/library/scan.html#simple-accumulation-into-a-scalar-ditching-lambda) I used the following:

```auto
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:

```auto
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.

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [May 5, 2023, 8:13am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/9 "2023-05-05T08:13:35Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [May 5, 2023, 8:15am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/10 "2023-05-05T08:15:49Z")

</div>

So can I just do:

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

```

Or is the syntax different?

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [May 5, 2023, 8:20am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/11 "2023-05-05T08:20:41Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [May 5, 2023, 8:37am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/12 "2023-05-05T08:37:35Z")

</div>

Yes that works, but I need to use:

```auto
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.

---

<div class="post-metadata">

**Author:** ![M\_K-C](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/m_k-c/32/6150_2.png) [@M\_K-C](https://discourse.pymc.io/u/M_K-C)\
**Post date:** [May 31, 2023, 4:13pm UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/13 "2023-05-31T16:13:02Z")

</div>

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` ?

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [June 1, 2023, 1:09am UTC](https://discourse.pymc.io/t/using-pytensor-scan-to-loop-over-multidimensional-array/11976/14 "2023-06-01T01:09:22Z")

</div>

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

```auto
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.