From Sympy to aesara.scan

[FabiJa]

Hi,
I am trying to export Sympy expressions to aesara. Theses sympy expressions have to be evaluated in a recursive fashion and I want to realize this with aesara.scan. Of course I can use sympy.printing.aesaracode.aesara_function for this and evaluate the resulting aesara function in a Python loop. But maybe scan gives better performance?

Here a MWE
Variant 1 with sympy.printing.aesaracode.aesara_function

import aesara
import numpy as np
import aesara.tensor as at
from sympy import symbols
from sympy.printing.aesaracode import aesara_function, aesara_code

# set-up the sympy expressions
x, y, z = symbols('x y z')
z = x + 2*y
rslt = 3*z

rslt_aesara = aesara_function(inputs=[x, y], outputs=[rslt])

n_steps = int(1e6)  # steps of recursion
x = np.zeros([n_steps])
y = np.linspace(1.0, 100.0, num=n_steps)

for idx in range(n_steps-1):  # recursion
    x[idx+1] += rslt_aesara(x[idx], y[idx])

rslt_python = x  # result of recursion for the mixed aesara/Python variant

Variant 2 with native aesara and scan

# pure aesara variant with scan
x_ = at.dvector()
y_ = at.dvector()

# define wrapper function because `scan` cannot operate on a compuational graph directly?
def rslt_wrapper(y, x):
    z = x + 2*y
    return 3*z

# use `scan` for the recursion
output, updates = aesara.scan(fn=rslt_wrapper, sequences=[y_,], outputs_info=[x_], n_steps=n_steps-1)

# compile aesara graph
solution = aesara.function(inputs=[x_, y_],
                    outputs=output,
                    updates=updates)
rslt_aesara = np.zeros([n_steps])
rslt_aesara[1:] = solution([x[0]], y)[:,0]

rslt_aesara  # result of recursion for the pure aesara variant

So, how can I combine the two variants, where I start with the sympy expressions and then use scan for the recursion?

Thx.

[brandonwillard]

You're touching on some of the most complicated parts of Aesara here!

You could construct a Scan somewhat manually using the underlying graph constructed by aesara_function(inputs=[x, y], outputs=[rslt]) (i.e. from SymPy) as the loop body.

You might even be able to do this straight from within the aesara.scan fn (i.e. your rslt_wrapper), or by creating an OpFromGraph for the SymPy-converted operations and returning that from within the fn.

[FabiJa]

Your first option, i.e. using aesara_function and evaluate this in a Python loop, was my first approach, where I started (see variant 1 above).

Could you elaborate a bit more on your option two?

Thx.

[brandonwillard]

What I mean is that there's a graph being produced by aesara.function that you could theoretically return in rslt_wrapper, but that's probably a bit too low-level for this case.

Instead, try returning the Aesara Variable results of aesara_function inside of rslt_wrapper. In other words, don't let aesara_function compile anything, just get the Aesara graph that it compiles. You'll need to look at the source for that.

[brandonwillard]

The toutputs variable from here is what you want to return in the scan fn (e.g. rslts_wrapper) function. You'll also need to make sure that tinputs are the arguments to the fn function. I believe you can map SymPy variables to existing Aesara variables via the cache argument used here.

It's starting to look like you only need to call aesara_code from within the fn function.