Looping over two dimensions with Pytensor Scan (original) (raw)

Hi,

I would like to use pytensor.scan to ‘loop’ over two dimensions of a matrix separately. How can this best be done?
I seem to run into errors when I try to do it, but I also don’t fully understand how scan works.

Thank you for your help!

Can you show what you’ve tried so far?

M_K-C May 27, 2024, 4:36pm 3

Yes,
for now I was just trying to loop over a function and store the result in a separate variable:

def function(x,y):
    return 2*x + y

def multi_loop():
    result_inner, _ = pytensor.scan(fn=function, sequences=[x], non_sequences=[y])
    y = result_inner[-1]
    return y

result_outer, _ = pytensor.scan(fn=multi_loop, sequences = [pt.arange(10)])

result_final = result_outer

This is a simplified version to show what I am trying to do. When running this I get an error saying TypeError: TensorType does not support iteration. Maybe you are using builtins.sum instead of pytensor.tensor.math.sum? (Maybe .max?)

I think the issue is how I have set up the outer loop. The inner one runs fine on its own.
Any help is highly appreciated

Shouldn’t your outer sequences be 2-dimensional?

M_K-C May 27, 2024, 5:36pm 5

Yes, true, but even then I get the same error

Your code as posted doesn’t raise any errors, though I had to make some small changes to your functions:

import pytensor 
import pytensor.tensor as pt

def function(x,y):
    return 2*x + y

def multi_loop(x, y):
    result_inner, _ = pytensor.scan(fn=function, sequences=[x], non_sequences=[y])
    y = result_inner[-1]
    return y

x,y = pt.matrices('x', 'y')
result_outer, _ = pytensor.scan(fn=multi_loop, non_sequences=[x, y], n_steps=10)
result_outer.eval({x:np.eye(3), y:np.eye(3)})

Can you post a complete example that raises the error you’re encountering?