BUG: Scan inner graphs are not optimized in NUMBA / JAX backends

[ricardoV94]

Reproducable code example:

import aesara
aesara.config.mode = "NUMBA"  # Otherwise it works
import aesara.tensor as at
from aesara.compile.builders import OpFromGraph

x = at.scalar("x")
out = at.log(x)
op = OpFromGraph([x], [out], inline=True)

xs = at.vector("xs")
seq, _ = aesara.scan(
    fn=lambda x: op(x),
    sequences=[xs],
)

f = aesara.function([xs], seq)
aesara.dprint(f)

for{cpu,scan_fn} [id A] 5
 |Shape_i{0} [id B] 0
 | |xs [id C]
 |Subtensor{int64:int64:int8} [id D] 4
 | |xs [id C]
 | |ScalarFromTensor [id E] 3
 | | |Elemwise{Composite{Switch(LE(i0, i1), i1, i2)}} [id F] 2
 | |   |Shape_i{0} [id B] 0
 | |   |TensorConstant{0} [id G]
 | |   |TensorConstant{0} [id H]
 | |ScalarFromTensor [id I] 1
 | | |Shape_i{0} [id B] 0
 | |ScalarConstant{1} [id J]
 |Shape_i{0} [id B] 0

Inner graphs:

for{cpu,scan_fn} [id A]
 >OpFromGraph{inline=True} [id K]
 > |*0-<TensorType(float64, ())> [id L] -> [id D]

OpFromGraph{inline=True} [id K]
 >Elemwise{log,no_inplace} [id M]
 > |*0-<TensorType(float64, ())> [id L]

[ricardoV94]

I also can't figure out how to make it work in Numba mode, for some reason it doesn't apply then

[ricardoV94]

Some progress, it seems that when in NUMBA mode (and I assume JAX as well), the inner fgraph of scan is not optimized according to the Scan mode. In the default backend the optimization is triggered when accessing the dynamic property node.fn, which seems to usually be done first in make_thunk. I imagine this path is not triggered by the JIT linkers.

If I access the property manually, the correct graph is obtained:

import aesara
aesara.config.mode = "NUMBA"
import aesara.tensor as at
from aesara.compile.builders import OpFromGraph

x = at.scalar("x")
out = at.log(x)
op = OpFromGraph([x], [out], inline=True)

xs = at.vector("xs")
seq, _ = aesara.scan(
    fn=lambda x: op(x),
    sequences=[xs],
)

seq.owner.op.fn  # Trigger inner fgraph optimization manually!!!

# aesara.config.optimizer_verbose = True
f = aesara.function([xs], seq)
aesara.dprint(f)

[ricardoV94]

#28 solves this for the NUMBA backend, but not JAX. I did not try to change it since Scan is completely broken for JAX