From Sympy to aesara.scan #1528

FabiJa · 2024-04-12T07:50:43Z

FabiJa
Apr 12, 2024

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 · 2024-04-15T21:04:05Z

brandonwillard
Apr 15, 2024
Maintainer

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.

0 replies

FabiJa · 2024-04-17T07:04:39Z

FabiJa
Apr 17, 2024
Author

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.

2 replies

brandonwillard Apr 17, 2024
Maintainer

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 Apr 17, 2024
Maintainer

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.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

From Sympy to aesara.scan #1528

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 2 comments 2 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Select a reply

Uh oh!

Uh oh!

From Sympy to aesara.scan #1528

Uh oh!

Uh oh!

FabiJa Apr 12, 2024

Replies: 2 comments · 2 replies

Uh oh!

brandonwillard Apr 15, 2024 Maintainer

Uh oh!

FabiJa Apr 17, 2024 Author

Uh oh!

Uh oh!

brandonwillard Apr 17, 2024 Maintainer

Uh oh!

Uh oh!

brandonwillard Apr 17, 2024 Maintainer

FabiJa
Apr 12, 2024

Replies: 2 comments 2 replies

brandonwillard
Apr 15, 2024
Maintainer

FabiJa
Apr 17, 2024
Author

brandonwillard Apr 17, 2024
Maintainer

brandonwillard Apr 17, 2024
Maintainer