Using PyTensor.scan to loop over multidimensional array

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