aesara
aesara copied to clipboard
Constant folding does not obey some evaluation expectations
The way that aesara.tensor.basic_opt.constant_folding
works does not obey the evaluation expectations of some Op
s like IfElse
and the stability expectations provided by our stabilizing rewrites.
In general, constant_folding
is a local rewrite that works from the inputs to the outputs of a graph. Because of this, it violates the "lazy" evaluation expectations of IfElse
(i.e. it computes both branches). See the conversation here.
Also, because constant_folding
is applied before our stabilizing rewrites, expressions that would yield finite/non-NaN values after stabilization will not when constant folded. See https://github.com/aesara-devs/aesara/issues/1072.
Both of these issues are distinct, but they point to the need for us to reconsider how constant folding is performed.
Here's an example test method that can be added to tests.test_ifelse.TestIfElse
:
def test_constant_folding(self):
class TrackingOp(Op):
n_evaled = 0
def make_node(self, inp):
return Apply(self, [inp], [inp.type()])
def perform(self, node, inputs, output_storage):
self.n_evaled += 1
output_storage[0][0] = inputs[0]
a = TrackingOp()(at.as_tensor(0))
b = TrackingOp()(at.as_tensor(1))
c = at.as_tensor(1)
from aesara.link.vm import VMLinker
# Use the Python-based `Stack` so we can walk through its lazy evaluation (for debugging/inspection)
linker = VMLinker(use_cloop=False, allow_partial_eval=True)
f = function([], IfElse(1)(c, a, b), mode=Mode(linker=linker, optimizer=[]))
# .including("topo_constant_folding"))
assert a.owner.op.n_evaled == 1
assert b.owner.op.n_evaled == 0