Automatically "metatize" backend helper functions

[brandonwillard]

It would be great if we could automatically transform backend helper/graph constructor functions into meta equivalents.

The problem

These "helper functions" are the standard Python functions found in the theano.tensor and tensorflow modules (e.g. theano.tensor.eye/tensorflow.linalg.eye, etc.) . They're generally used to simplify the construction of graphs, which is actually done by creating Apply/Operation objects—as well as the objects underlying those (e.g. OpDefs and NodeDefs for TensorFlow and sometimes Ops in Theano).

When we want to craft a meta graph—from the ground up—corresponding to the kind of graph that one of these helper functions would produce (e.g. in the normal course of using TensorFlow or Theano), we unfortunately have to do this at the underlying Op/OpDef level and effectively reproduce the code within these helper functions. In other words, we would like to have meta versions of these helper functions (e.g. mt.eye as the meta version of tensorflow.eye) that more-or-less do the same things as the originals, but use meta Ops/OpDefs instead, and—ideally—they would even work with logic variable inputs (e.g. an mt.eye that's given a meta tensor with logic variables for shape and/or dtype).

For instance, tensorflow.abs is one such function; it takes a tensor/graph as input and does a simple check to determine which OpDef (i.e. Abs or ComplexAbs) to use when it constructs the output graph representing an absolute value of the input. That check is simply a condition on the dtype.is_complex property of the input tensor.

The steps in tensorflow.abs can largely be applied to symbolic_pymc's meta tensors without much/any changes. Even so, tensorflow.abs will necessarily construct TensorFlow objects and not the corresponding meta ones we actually want.

The question is: how do we [re]use as much of the existing helper function code as possible without entirely rewriting them? Of course, it could be quite an undertaking to cover every case, but there might be a few cheap work-arounds that help in more than a few cases.

FYI: This applies to both Theano and TensorFlow backends.

An example AST-based approach

In Theano, tt.diagonal is a plain function and won't produce meta objects or accept them as arguments. However, the implementation of tt.diagonal is extremely simple: it just constructs an ExtractDiag operator instance. Since we can create a meta version of the ExtractDiag operator, we just need tt.diagonal to use it.

The following demonstrates how we could automatically convert some simple helper functions into meta function equivalents using straight-forward AST manipulation.

import ast
import astor
import inspect
import types
import theano.tensor as tt

from symbolic_pymc.theano.meta import mt


class RewriteName(ast.NodeTransformer):
    def __init__(self):
        self.original_name = None
        self.new_name = None

    def visit_FunctionDef(self, node):
        """This is the first/outer-most function that we're attempting to
        change.
        """
        if self.original_name and node.name == self.original_name:
            # There's a function definition that shadows the outer
            # function.  That means we shouldn't rename calls to
            # the original/outer function, since they're actually
            # for this inner function.
            self.new_name = self.original_name
        elif self.new_name is None and not node.name.startswith('mt_'):
            self.original_name = node.name
            self.new_name = f'mt_{node.name}'
            node.name = self.new_name
        return self.generic_visit(node)

    def visit_Call(self, node):
        if getattr(node.func, 'id', None) == self.original_name:
            node.func.id = self.new_name
        return self.generic_visit(node)

    def visit_Name(self, node):
        new_node = node
        node_obj = vars(tt).get(node.id)
        if node_obj:
            try:
                # This will throw if there's no meta object (or one cannot be
                # created).
                mt(node_obj)
            except ValueError:
                print(f'Meta object not found for {node_obj}')
            else:
                # The meta object exists, so go ahead with the change.
                new_node = ast.copy_location(
                    ast.Attribute(value=ast.Name(id='mt', ctx=ast.Load()),
                                  attr=node.id,
                                  ctx=node.ctx),
                    node)
        return new_node


tt_diagonal_src = inspect.getsource(tt.diagonal)

Here's the source for the original Theano function we want to make compatible with our meta objects:

>>> print(tt_diagonal_src)
def diagonal(a, offset=0, axis1=0, axis2=1):
    """
    A helper function for `theano.tensor.ExtractDiag`. It accepts tensor with
    `ndim >= 2` as input. The name `diagonal` is just meant to keep it
    consistent with numpy.

    Parameters
    ----------
    a : symbolic tensor
    offset : int
        offset
    axis1 : int
    axis2 : int

    Returns
    -------
    tensor : symbolic tensor

    """
    return ExtractDiag(offset, axis1, axis2)(a)

Now, we run the AST transform:

diagonal_ast = ast.parse(tt_diagonal_src)
new_diagonal = RewriteName().generic_visit(diagonal_ast)

and view the [source for the] transformed Theano helper function:

>>> print(astor.to_source(new_diagonal))
def mt_diagonal(a, offset=0, axis1=0, axis2=1):
    """
    A helper function for `theano.tensor.ExtractDiag`. It accepts tensor with
    `ndim >= 2` as input. The name `diagonal` is just meant to keep it
    consistent with numpy.

    Parameters
    ----------
    a : symbolic tensor
    offset : int
        offset
    axis1 : int
    axis2 : int

    Returns
    -------
    tensor : symbolic tensor

    """
    return mt.ExtractDiag(offset, axis1, axis2)(a)

Simply put, we've automatically made the change from Theano's ExtractDiag to our own meta ExtractDiag (via the meta accessor mt).

The following will create the transformed function in the current namespace/scope:

mt_diagonal_code = compile(ast.fix_missing_locations(new_diagonal), '<meta>',
                           mode='exec')
exec(mt_diagonal_code)

One major shortcoming to this approach involves how the converted meta objects are used. For instance, some conditions in these helper functions involve comparisons that aren't sound when performed with meta objects (e.g. inequalities involving fields populated by logic variables). However, in this case, it's possible that large sets of such restrictions could be lifted by implementing "typed" logic variables (e.g. logic variables in numeric/array-valued fields that implement numeric comparisons, etc.)