Construct graphs resulting from arbitrary applications of rewrites

[brandonwillard]

• Introduction
  • Generalized Apply Nodes
  • Traversing the Space of Rewrites

We need the ability to efficiently construct the graphs resulting from arbitrarily ordered applications of individual rewrites.

This functionality is required for our work in AeMCMC, which after aesara-devs/aemcmc#45 has moved the rewrite considerations into the realm of applied vs. unapplied rewrites. More specifically, AeMCMC needs to "temporarily" consider scale-mixture expanded random variables (see aesara-devs/aemcmc#46) as a means of discovering custom Gibbs samplers. Those expansions need to be ephemeral, because some may end up being pointless–or potentially deleterious–when they don't serve to produce a sampler.

Consider the following two rewrites:

from collections import Counter

import aesara
import aesara.tensor as at
from aesara.graph.opt import local_optimizer
from aesara.scalar.basic import Exp, Log


@local_optimizer([at.log])
def local_log_of_exp(fgraph, node):
    """Rewrite ``at.log(at.exp(x))`` to ``x``."""
    x = node.inputs[0]

    if isinstance(x.owner.op.scalar_op, Exp) and isinstance(node.op.scalar_op, Log):
        new_out = x.owner.inputs[0]
        return [new_out]


@local_optimizer([at.add])
def local_reduce_add(fgraph, node):
    """Rewrite ``x + x`` to ``2 * x``."""

    unique_inputs = Counter(node.inputs)

    if all(v == 1 for v in unique_inputs.values()):
        return None

    if len(unique_inputs) > 1:
        res = [at.add(*(v * k if v > 1 else k for k, v in unique_inputs.items()))]
    else:
        res = [v * k for k, v in unique_inputs.items()]

    return res

and the following example graph:

x = at.vector("x")
z = at.log(at.exp(x + x)) + x + x

aesara.dprint(z)
# Elemwise{add,no_inplace} [id A]
#  |Elemwise{add,no_inplace} [id B]
#  | |Elemwise{log,no_inplace} [id C]
#  | | |Elemwise{exp,no_inplace} [id D]
#  | |   |Elemwise{add,no_inplace} [id E]
#  | |     |x [id F]
#  | |     |x [id F]
#  | |x [id F]
#  |x [id F]

In other words, we're rewriting the expression $\log\exp\left(x + x\right) + x + x$ and we want to record all the possible graphs resulting from iteratively applying each rewrite in different orders.

We can anticipate rewrites like the following:

$$ \begin{align} \log\exp\left(x + x\right) + x + x &\to \log\exp\left(2 x\right) + x + x &\to 2 x + x + x &\to 2 x + 2 x \\ &\to \log\exp\left(x + x\right) + 2 x &\to x + x + 2 x &\to 2 x + 2 x \\ &\to x + x + x + x &\to 4 x & \\ \vdots &\vdots &\vdots &\vdots \\ \end{align} $$

If we want to apply the above two rewrites to a graph, we must first choose a means of traversing the graph and applying them to each node. In this case, nearly any topological traversal order will do (as long as it visits every node), but the traversal processes must be performed repeatedly so that the graphs resulting from each successfully applied rewrite can themselves be rewritten.

The EquilibriumOptimizer is currently the most suitable Rewriter for this kind of job. It will construct a LocalOptGroup and apply the two rewrites until a fixed point is reached.

N.B. LocalOptGroup uses a "tracking" mechanism to make sure that the above two rewrites are only attempted when a given node uses an at.log or at.add Op–i.e. the two Ops registered by each rewrite in their local_optimizer decorators.

As with all rewrites, we need to construct a FunctionGraph object and apply a Rewriter to it:

from aesara.graph.opt import EquilibriumOptimizer
from aesara.graph.fg import FunctionGraph

from aesara.tensor.math_opt import local_add_mul_fusion


fg = FunctionGraph(outputs=[z])

eq_rewriter = EquilibriumOptimizer(
    [
        local_reduce_add, local_log_of_exp,
        # We need this to turn `at.add(a, at.add(b, c))` into `at.add(a, b, c)`
        local_optimizer([at.add])(local_add_mul_fusion)
     ],
    max_use_ratio=aesara.config.optdb__max_use_ratio,
)

_ = eq_rewriter.apply(fg)

aesara.dprint(fg)
# Elemwise{add,no_inplace} [id A] 4
#  |Elemwise{mul,no_inplace} [id B] 3
#  | |InplaceDimShuffle{x} [id C] 2
#  | | |TensorConstant{2} [id D]
#  | |x [id E]
#  |Elemwise{mul,no_inplace} [id F] 1
#    |InplaceDimShuffle{x} [id G] 0
#    | |TensorConstant{2} [id H]
#    |x [id E]

As we know, only one of the anticipated sequences of rewrites has been applied. If we want to observe all sequences, we need to alter the steps applied by EquilibriumOptimizer so that a new FunctionGraph instance is produced each time a rewrite is successfully applied, then a queue/"stream" of FunctionGraph instances can be constructed and processed by each rewrite in varying orders.

First, why are new FunctionGraph instances needed? Because our rewriting interfaces are all based on FunctionGraphs, and information like FunctionGraph.clients and other data structures are maintained by Features attached to a FunctionGraph. Information in the attached Features is updated incrementally through callbacks triggered by a FunctionGraph when nodes are added, removed, and updated.

Since the incrementally built data structures in a FunctionGraph and its Features are used all throughout our rewrite implementations, we need them to work in this new situation where independent graphs are repeatedly produced by rewrites.

At first, it might seem like we can simply clone FunctionGraphs at a certain point; however, since FunctionGraph.replace works by altering Apply nodes in-place, every single Apply node in a graph must be cloned, and this invalidates much of the information tracked by Features and the like, because that information often consists of dicts mapping Apply nodes and/or their output Variables.

One might notice that the situation can be reformulated into FunctionGraph "diffs" (i.e. each rewrite applies changes to a graph via FunctionGraph.replace and those replacement pairs/maps can be collected and applied to a reference/initial FunctionGraph. Functionality based on this reformulation is provided by the History Feature, which "snapshots" a FunctionGraph instance, collects the changes/diffs, and can apply them in "reverse" to revert a FunctionGraph to a previous state. Because History uses FunctionGraph.change_node_input, all the listening Features can update their internal data structures accordingly.

This approach still doesn't address the main issue: Apply nodes are updated in-place, making it necessary to clone entire FunctionGraphs at some point.

N.B. #666 addresses some of these cloning issues, but it still needs to recurse into the clients of a cloned Apply node and clone the client nodes. In other words, any time an Apply node's inputs are changed, that node needs to be cloned, as well as all the Apply nodes that depend on the updated node's outputs.

Using our example graph z, i.e.

# Elemwise{add,no_inplace} [id A]
#  |Elemwise{add,no_inplace} [id B]
#  | |Elemwise{log,no_inplace} [id C]
#  | | |Elemwise{exp,no_inplace} [id D]
#  | |   |Elemwise{add,no_inplace} [id E]
#  | |     |x [id F]
#  | |     |x [id F]
#  | |x [id F]
#  |x [id F]

if we want to replace the node/Variable labeled E (as we would with local_reduce_add) and produce a new, independent graph, we need to clone the "downstream" nodes (i.e. D, C, B, and A), because E is an input to all of them.

This cloning simplifies some things, but it also demands resources and doesn't scale well with the size of a graph. The same goes for all the information collected about the cloned Apply nodes; it all needs to be recomputed. These recomputations can be very redundant as well. Just consider the FunctionGraph.clients information. The only client of node B is A, and that relationship doesn't change when A and B are cloned; nevertheless, the entry for B's clients in FunctionGraph.clients needs to be recomputed using the cloned values.

So how can we avoid the need to clone downstream Apply nodes and invalidate collected information?

Generalized Apply Nodes

One way to look at this cloning situation: Apply nodes can only have one set of inputs, which limits their reuse.

Recall that Apply nodes are simply containers that hold an Op and lists of input and output Variables. If we allow Apply nodes to have multiple sets of inputs, then the example replacement mentioned above would only involve the addition of an alternative input to node D.

Using this approach an Apply node is now determined by its Op, output Variable(s), and input Types–instead of input Variables. A "context" is needed to determine the input Variable(s), and, if we connect these contexts to FunctionGraphs, we get a little closer to our goal.

Here's a MWE illustrating these alternative Apply nodes:

from contextlib import contextmanager

import aesara
import aesara.tensor as at
from aesara.graph.basic import Apply
from aesara.graph.type import Type
from aesara.graph.utils import Scratchpad
from aesara.graph.fg import FunctionGraph

__graph_context__ = None


@contextmanager
def set_input_context(context):
    global __graph_context__
    last_context = __graph_context__
    __graph_context__ = context
    try:
        yield
    finally:
        __graph_context__ = last_context


class ContextualApply(Apply):
    def __init__(self, op, in_types, outputs):

        inputs_are_types = all(isinstance(in_t, Type) for in_t in in_types)

        if not inputs_are_types:
            if __graph_context__ is None:
                raise ValueError("No context set for ContextualApply nodes")

            self.input_types = tuple(inp.type for inp in in_types)
        else:
            self.input_types = in_types

        self.op = op
        self.outputs = []
        self.tag = Scratchpad()
        for idx, out in enumerate(outputs):
            out.owner = self
            out.index = idx
            self.outputs.append(out)

        if __graph_context__ is not None:
            __graph_context__[self] = list(in_types)

    @property
    def inputs(self):
        if __graph_context__ is None:
            raise ValueError("No context set for ContextualApply nodes")

        return __graph_context__[self]

    @inputs.setter
    def inputs(self, val):
        if __graph_context__ is None:
            raise ValueError("No context set for ContextualApply nodes")

        if not all(v.type == in_type for v, in_type in zip(val, self.input_types)):
            raise ValueError(f"Inputs must match types {self.input_types}")

        __graph_context__[self] = val

    def __str__(self):
        return repr(self)

    def __repr__(self):
        if __graph_context__ is None:
            _inputs = self.input_types
        else:
            _inputs = __graph_context__[self]

        return f"{type(self).__name__}({self.op}, {_inputs}, {self.outputs})"


in_1, in_2, in_3, in_4 = at.vectors(*[f"in_{i}" for i in "1234"])
out_1, out_2 = at.vectors("out_1", "out_2")

context_1 = {}

with set_input_context(context_1):
    node_1 = ContextualApply(at.add, [in_1, in_2], [out_1])
    node_2 = ContextualApply(at.add, [out_1, in_3], [out_2])

# Printing an `Apply` node outside of a context shows its input types
print(node_1)
# ContextualApply(Elemwise{add,no_inplace}, (TensorType(float64, (None,)), TensorType(float64, (None,))), [out_1])

# Printing inside of a context shows its input `Variable`s

with set_input_context(context_1):
    print(node_1)
# ContextualApply(Elemwise{add,no_inplace}, [in_1, in_2], [out_1])

# We create a cloned context
context_2 = {k: list(v) for k, v in context_1.items()}

# and change an input within the new context
with set_input_context(context_2):
    node_1.inputs[1] = in_4
    print(node_1)
# ContextualApply(Elemwise{add,no_inplace}, [in_1, in_4], [out_1])

# The original context is unchanged, by design

with set_input_context(context_1):
    print(node_1)
# ContextualApply(Elemwise{add,no_inplace}, [in_1, in_2], [out_1])

# Most/all functionality should work as usual within a given graph context

with set_input_context(context_1):
    fg_1 = FunctionGraph(outputs=[out_2], clone=False)
    aesara.dprint(fg_1)
# Elemwise{add,no_inplace} [id A] 'out_2' 1
#  |Elemwise{add,no_inplace} [id B] 'out_1' 0
#  | |in_1 [id C]
#  | |in_2 [id D]
#  |in_3 [id E]


with set_input_context(context_2):
    fg_2 = FunctionGraph(outputs=[out_2], clone=False)
    aesara.dprint(fg_2)
# Elemwise{add,no_inplace} [id A] 'out_2' 1
#  |Elemwise{add,no_inplace} [id B] 'out_1' 0
#  | |in_1 [id C]
#  | |in_4 [id D]
#  |in_3 [id E]


# The collected information in each `FunctionGraph` is context specific

fg_1.clients[in_2]
# [(ContextualApply(Elemwise{add,no_inplace}, (TensorType(float64, (None,)), TensorType(float64, (None,))), [out_1]),
#   1)]
fg_2.clients[in_4]
# [(ContextualApply(Elemwise{add,no_inplace}, (TensorType(float64, (None,)), TensorType(float64, (None,))), [out_1]),
#   1)]
in_2 not in fg_2.clients
# True


out_3 = at.vector("out_3")

with set_input_context(context_2):
    node_3 = ContextualApply(at.mul, [in_1, in_2], [out_3])
    fg_2.replace(in_3, out_3, import_missing=True)
    aesara.dprint(fg_2)
# Elemwise{add,no_inplace} [id A] 'out_2' 2
#  |Elemwise{add,no_inplace} [id B] 'out_1' 1
#  | |in_1 [id C]
#  | |in_4 [id D]
#  |Elemwise{mul,no_inplace} [id E] 'out_3' 0
#    |in_1 [id C]
#    |in_2 [id F]

These contexts still require that one clone FunctionGraphs, but at least they make it possible to do so without cloning most of the Variables and Apply nodes involved. Instead, it might be possible to get away with only shallow clones of FunctionGraphs and the data structures in Features that track node-related information.

Traversing the Space of Rewrites

Concerns similar to the ones above arise when we consider rewrites that require the use of associative-commutative (AC) properties, because we don't–for instance–want to permute the order of arguments to a commutative Op, construct the new Apply node, make the change in a FunctionGraph, call all the listening Features, let them perform their updates, and then see if that ordering leads to a graph that matches the next step in a rewrite. Again, in this sort of scenario, we need to temporarily manifest these graphs in a lower-cost way and then reason about them.

In #634, AC properties are implemented using miniKanren (via kanren). miniKanren allows one to easily traverse and evaluate a few–or all–graphs resulting from every application order of each rewrite.

Simply put, miniKanren provides a framework for lazily producing and processing sequences of rewrites applied in every order. These sequences can be infinite, as well.

Let's take a look at what miniKanren does with our graphs.

from etuples import etuplize
from kanren import run, conso, vars
import aesara.tensor as at
from cons import cons, car, cdr


x = at.vector("x")
z = at.log(at.exp(x + x)) + x + x

# First, we need to convert our graphs into `etuple` form, which is like
# converting `Apply` nodes into tuples--at least the `Op` and inputs parts; the
# outputs are implied as the result of evaluation.
# We do this because all of the standard miniKanren relational goals are designed
# for `cons` semantics, which can be most easily and directly mapped to Python sequence types, so
# converting our graphs to nested sequences makes them instantly amenable to most goals.
z_et = etuplize(z)
# e(
#   e(
#     aesara.tensor.elemwise.Elemwise,
#     <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#     <frozendict {}>),
#   e(
#     e(
#       aesara.tensor.elemwise.Elemwise,
#       <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#       <frozendict {}>),
#     e(
#       e(
#         aesara.tensor.elemwise.Elemwise,
#         <aesara.scalar.basic.Log at 0x7f32d63094c0>,
#         <frozendict {}>),
#       e(
#         e(
#           aesara.tensor.elemwise.Elemwise,
#           <aesara.scalar.basic.Exp at 0x7f32d6309700>,
#           <frozendict {}>),
#         e(
#           e(
#             aesara.tensor.elemwise.Elemwise,
#             <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#             <frozendict {}>),
#           x,
#           x))),
#     x),
#   x)

# To illustrate the `cons` semantics, lets take a look at the fundamental functions:
car(z_et)
# e(
#   aesara.tensor.elemwise.Elemwise,
#   <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#   <frozendict {}>)
cdr(z_et)
# e(
#   e(
#     e(
#       aesara.tensor.elemwise.Elemwise,
#       <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#       <frozendict {}>),
#     e(
#       e(
#         aesara.tensor.elemwise.Elemwise,
#         <aesara.scalar.basic.Log at 0x7f32d63094c0>,
#         <frozendict {}>),
#       e(
#         e(
#           aesara.tensor.elemwise.Elemwise,
#           <aesara.scalar.basic.Exp at 0x7f32d6309700>,
#           <frozendict {}>),
#         e(
#           e(
#             aesara.tensor.elemwise.Elemwise,
#             <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#             <frozendict {}>),
#           x,
#           x))),
#     x),
#   x)
z_et == cons(car(z_et), cdr(z_et))
# True

# Create two logic variables to be unified with terms in and/or produced by the
# relational goals.
# These serve as the "unknowns" in the "equations" we want to solve.
car_lv, cdr_lv = vars(2)

res = run(
    # This is the number of results we want.  `0` means that we want a
    # tuple with all the results.
    0,
    # These are the results we want reified (i.e. the logic variables replaced
    # with their unified values)
    (car_lv, cdr_lv),
    # These are the relational goals that must be fulfilled.  They ultimately
    # determine the unified values of each logic variable.
    # In this case, we use the `conso` goal, which implements the relation
    # `cons(car(a), cdr(b)) == c`.  It is equivalent to
    # `eq(cons(car(a), cdr(b)), c)`.
    conso(car_lv, cdr_lv, z_et)
)

# The first (and only) result satisfying the goals
res[0]

# The reified value of `car_lv` in the first result
res[0][0] == car(z_et)
# True
res[0][1] == cdr(z_et)
# True

miniKanren works by producing branching streams of "state" objects. The states in miniKanren are usually unification states: i.e. maps of logic variables and their unified values. As miniKanren process each goal, those goals add new states to the stream–states which satisfy their goals, of course. When a goal-satisfying state cannot be produced in a given stream, that stream ends, yielding no results.

Let's take a look at the internal states produced by kanren.

from kanren.core import lall

# First, this is what the string forms of our logic variables look like:
car_lv, cdr_lv
# (~_5, ~_6)

# We evaluate the same goal as above using the stream
initial_state = {}
next_state = next(lall(conso(car_lv, cdr_lv, z_et))(initial_state))
next_state
# {~_5: e(
#    aesara.tensor.elemwise.Elemwise,
#    <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#    <frozendict {}>),
#  ~_6: e(
#    e(
#      e(
#        aesara.tensor.elemwise.Elemwise,
#        <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#        <frozendict {}>),
#      e(
#        e(
#          aesara.tensor.elemwise.Elemwise,
#          <aesara.scalar.basic.Log at 0x7f32d63094c0>,
#          <frozendict {}>),
#        e(
#          e(
#            aesara.tensor.elemwise.Elemwise,
#            <aesara.scalar.basic.Exp at 0x7f32d6309700>,
#            <frozendict {}>),
#          e(
#            e(
#              aesara.tensor.elemwise.Elemwise,
#              <aesara.scalar.basic.Add at 0x7f32d62f45e0>,
#              <frozendict {}>),
#            x,
#            x))),
#      x),
#    x)}

initial_state = {car_lv: 1}
next(lall(conso(car_lv, cdr_lv, z_et))(initial_state))
# StopIteration

The second state example demonstrates a failure of the state to satisfy the relation implemented by the goal (i.e. cons(1, cdr(cdr_lv)) != z_et for any cdr_lv).

In Aesara, graphs are walked via recursive calls to nodes' inputs. In kanren, nodes are effectively etuples, and their inputs are obtained via a call to cdr, making recursive conso goals a relational means of walking a graph. See here for detailed illustrations of relational graph walking in kanren.

The miniKanren states produced when graphs are walked look similar to the example state printed above: they contain logic variables for the car and cdr values of each node in a graph.

from kanren import eq
from kanren.graph import walko


# A simple tuple-based graph that's easier to investigate than an Aesara graph
graph = (1, (2, (4, 5)))

# A logic variable representing the "output" of the relation that follows
res_lv = var()
res_lv
# ~_139

# Apply the goal `eq` to each matching element in the graphs between `graph`
# and `res_lv`.
# This goal implements a relation like `apply(eq, graph) == res_lv`.
goal = walko(eq, graph, res_lv)

initial_state = {}
stream = lall(goal)(initial_state)

# First node
next_state = next(stream)
next_state
# {~_139: (1, (2, (4, 5)))}

next_state = next(stream)
next_state
# {~_146: 1,
#  ~_147: ((2, (4, 5)),),
#  ~_139: ConsPair(~_144, ~_145),
#  ~_144: 1,
#  ~_158: (2, (4, 5)),
#  ~_159: (),
#  ~_145: ConsPair(~_156, ~_157),
#  ~_156: (2, (4, 5)),
#  ~_157: ()}

# After a few more iterations, we get the following:
next_state = next(stream)
next_state
# {~_146: 1,
#  ~_147: ((2, (4, 5)),),
#  ~_139: ConsPair(~_144, ~_145),
#  ~_144: 1,
#  ~_174: (2, (4, 5)),
#  ~_175: (),
#  ~_145: ConsPair(~_172, ~_173),
#  ~_182: 2,
#  ~_183: ((4, 5),),
#  ~_172: ConsPair(~_180, ~_181),
#  ~_180: 2,
#  ~_206: (4, 5),
#  ~_207: (),
#  ~_181: ConsPair(~_204, ~_205),
#  ~_204: (4, 5),
#  ~_205: (),
#  ~_173: ()}

Starting at the "output" logic variable res_lv, one can see that graph is being deconstructed into ConsPairs. ConsPairs are the base representation of two non-sequence objects that have been consed together (e.g. cons(1, 2) == ConsPair(1, 2)). Those ConsPairs hold logic variables created during recursive iterations within walko. Some of the logic variables map directly to the sub-graphs that haven't been traversed/"deconstructed" into conses yet. (N.B. a sequence represented in cons form is terminated by an empty sequence (e.g. cons(a, ()) == (a,)).)

These cons chains produced by walko are convenient representations for mixing and matching arbitrary parts of a graph. For example:

from unification import reify, var

# The original graph is encoded in the last state produced by miniKanren
reify(res_lv, next_state)
# (1, (2, (4, 5)))

# We can easily change any part of the graph by finding the appropriate logic
# variables and setting those values
new_state = next_state.copy()
new_state[var("_204")] = (6, 7)

reify(res_lv, new_state)
# (1, (2, (6, 7)))

new_state[var("_180")] = 3

reify(res_lv, new_state)
# (1, (3, (6, 7)))

The "multi-input" Apply node functionality illustrated in the section above can be reproduced using these unification states, because the "context" arrays in that framework are effectively encoded by the cons pair relationships in a unification state. This can be seen in the toy example above; by changing the value of a single logic variable, we were able to arbitrarily alter a nested tuple. If such a tuple corresponds to the inputs of a node in an Aesara graph, the resulting state corresponds to a new context array.

These unification states are considerably more general/flexible than the very Aesara-specific "multi-input" contexts, but they don't translate as directly to Aesara graphs and/or FunctionGraphs. Calling reify on a state will construct an Aesara graph almost from scratch (i.e. by walking the graph and creating creating etuples via conses and then evaluating those etuples to produce the Apply nodes). There is some caching in etuples, though, so a cons of parts that have already been consed before won't be re-evaluated in all instances.

In general, the graphs produced by miniKanren are not automatically reified at any point in the miniKanren process, so they don't ever take the form of standard Aesara graphs or FunctionGraph objects. This means it's not possible to use our existing Aesara rewrites and tooling on these values without costly reifications.

Nevertheless, it should be possible to combine the two techniques and leverage the streaming framework of miniKanren for our graph-space traversal needs.

To clarify, a kanren goal constructor and goal take the following general form:

from typing import MutableMapping, Generator


def relationo(*args):
    """Construct a goal for this relation."""

    def relationo_goal(S: MutableMapping) -> Generator[MutableMapping, MutableMapping, None]:
        """Generate states for the relation `relationo`.

        Parameters
        ----------
        S
            The miniKanren state (e.g. unification mappings/`dict`).

        Yields
        ------
        miniKanren states.

        """
        nonlocal args

        # Reify the terms of interest, which could be anything
        reified_vals = reify(..., S)

        # Create a new state
        S_new = S.copy()

        # Make changes to the state
        S_new[...] = ...

        # Add the new state to the stream
        yield S

    return relationo_goal

As the template above illustrates, we are able to do anything to the miniKanren state object (i.e. S in the example) within a goal. Also, we are able to use any type of state object that implements the MutableMapping interface.

Within this framework, we could sync the miniKanren states and their associated Aesara graphs and/or FunctionGraphs.

More on this later…