Using PyTensor.scan to loop over multidimensional array

Questions v5
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