Do we need to preserve undefined behavior in our rewrites?

[ricardoV94]

We have a rewrite that replaces exp(log(x)) -> x, but introduces a nan switch if the original value was negative.

import aesara
import aesara.tensor as at

x = at.scalar("x")
y = at.exp(at.log(x))
aesara.dprint(aesara.function([x], y))

Elemwise{Composite{Switch(GE(i0, i1), i0, i2)}} [id A] ''   0
 |x [id B]
 |TensorConstant{0} [id C]
 |TensorConstant{nan} [id D]

Do we need to do this, or should we assume the input values are valid for the sake of performance? For instance, we assume something like this when rewriting x/x->1

f = aesara.function([x], x/x)
print(f(0))  # 1.0
aesara.dprint(f)

DeepCopyOp [id A] ''   0
 |TensorConstant{1.0} [id B]

[brandonwillard]

We really should be conditioning rewrites like x / x -> 1 on their inputs; otherwise, the compiled functions aren't actually faithful to their original expressions.

If we want to avoid computing such conditions in a compiled graph, then we need to establish a means of specifying constraints on variables (e.g. x is a finite, non-zero real) so that the conditions on the rewrites can be discharged.

We would still need to validate such conditions by default when/if they're inputs to a compiled function, as we currently do via Types, but we could also provide the option to skip such checks. Regardless, the idea is that the conditions would no longer need to be in the graph.

It might be possible to extend our Types for this purpose, but we could also consider representing conditions in the graph. Something like the latter was done in Theano using a Hints type, I believe.

[ricardoV94]

I could not find any results for this Hint in Theano.

[brandonwillard]

See here.

More specifically, the hints I was alluding to are the ones like psd, which adds a positive semi-definite assumption to graph elements. Those are later used to perform specialized rewrites (e.g. psd_solve_with_chol).

I think this Hint approach was abandoned with the expectation that Variable.tag would serve the same purpose; however, Variable.tag is an inherently questionable approach (e.g. due to it adding mutable statefulness to Variables).

These features are also discussed in #499.

[brandonwillard]

You know, our Assert Op already plays a role similar to a Hint Op—at least when a Hint is used to specify a mathematical property via a predicate. In other words, z = aesara.assert_op(x, P(*a)) implies that the predicate P(*a) is true "for x", so, if P = lambda x, y: x > y, then using the output variable z adds the constraint x > y (i.e. a = (x, y)) to all the graphs containing z.
Really, any Op can serve this purpose by associating a property with its presence in a graph.

Use of an Op like Assert or Hint adds information to a graph in a "location" specific way, and that location information can be useful; however, it's also an annoyance for rewrites that need to match forms that don't care about or account for such information. The issue is easily solved by removing the respective nodes from the graph and storing the information they convey somewhere else.

The information conveyed by such Ops could—for example—be moved to the Variable.tags of their outputs via a simple rewrite that also removes the Ops/parent nodes.

The information could also be moved to the Python type level via Ops that output Variables with types and/or Types indicating the presence of a nullary property (i.e. a predicate with no arguments). Likewise, such Op nodes need only be removed from a graph at some point and replaced by owner-less clones of their original, specially type | Typeed outputs. It's possible that non-nullary properties (i.e. predicates with arguments) could be expressed at the type level through the construction of custom (sub)classes holding the inputs of the Apply nodes (i.e. the predicate arguments).

Either way, constraints/properties expressed as nodes in a graph can be removed and tracked through a few different mechanisms. The next set of questions/concerns involve the correspondence between constraint/property arguments (e.g. the x and y graphs in x > y) stored elsewhere (e.g. in Variable.tags or type | Types) and in the "originating" graph. Rewrites applied to such arguments present in the "originating" graph need to be reflected wherever else they're being tracked. This kind of correspondence is maintained through Variable.tag/type | Type cloning and in-place Apply updating, but it's not clear to me that this is truly sufficient.

Nevertheless, a Feature-based approach could always account for any insufficiencies by effectively tracking two different graphs: one with the Assert/Hint/Op-provided information and one without it. This starts to touch upon our discussions about "uncomputed" graph outputs and the like.

I'm somewhat partial to the type | Type approach, mostly because it allows us to easily leverage the Python type system (e.g. inheritance and formal, automated checking are useful). Regardless, in some sense, the Variable.tag approach is equivalent (e.g. store types in Variable.tag).

The basic idea is that we can use our existing graph framework to express constraints on terms, and we can reason about/with the constraints using our rewrite framework.