Description
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.)