Constant folding does not obey some evaluation expectations

[brandonwillard]

The way that aesara.tensor.basic_opt.constant_folding works does not obey the evaluation expectations of some Ops 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.

[brandonwillard]

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