Merge pull request #88 from brandonwillard/add-meta-operator

Refactor etuplization and introduce TensorFlow Operator meta class

Author: brandonwillard

symbolic_pymc/etuple.py

@@ -5,7 +5,7 @@
 
 from collections import Sequence
 
-from cons.core import ConsPair
+from cons.core import ConsPair, ConsNull
 
 from multipledispatch import dispatch
 
@@ -112,6 +112,8 @@ def eval_obj(self):
             evaled_kwargs = arg_grps.get(True, [])
 
             op = self._tuple[0]
+            op = getattr(op, "eval_obj", op)
+
             try:
                 op_sig = inspect.signature(op)
             except ValueError:
@@ -234,6 +236,8 @@ def etuplize(x, shallow=False, return_bad_args=False):
     """
     if isinstance(x, ExpressionTuple):
         return x
+    elif x is not None and isinstance(x, (ConsNull, ConsPair)):
+        return etuple(*x)
 
     try:
         # This can throw an `IndexError` if `x` is an empty
@@ -242,24 +246,22 @@ def etuplize(x, shallow=False, return_bad_args=False):
         args = arguments(x)
     except (IndexError, NotImplementedError):
         op = None
-        args = x
+        args = None
 
-    if not isinstance(args, ConsPair):
+    if not callable(op) or not isinstance(args, (ConsNull, ConsPair)):
         if return_bad_args:
             return x
         else:
             raise TypeError(f"x is neither a non-str Sequence nor term: {type(x)}")
 
-    # Not everything in a list/tuple should be considered an expression.
-    if not callable(op):
-        return etuple(*x)
-
     if shallow:
+        et_op = op
         et_args = args
     else:
+        et_op = etuplize(op, return_bad_args=True)
         et_args = tuple(etuplize(a, return_bad_args=True) for a in args)
 
-    res = etuple(op, *et_args, eval_obj=x)
+    res = etuple(et_op, *et_args, eval_obj=x)
     return res
 
 

symbolic_pymc/meta.py

@@ -239,6 +239,7 @@ def obj(self):
 
     @classmethod
     def base_subclasses(cls):
+        """Return all meta symbols with valid, implemented bases (i.e. base property is a type object)."""
         for subclass in cls.__subclasses__():
             yield from subclass.base_subclasses()
             if isinstance(subclass.base, type):
@@ -252,12 +253,18 @@ def __init__(self, obj=None):
         assert obj is None or isvar(obj) or isinstance(obj, self.base)
         self._obj = obj
 
+    @property
     def rands(self):
         """Get a tuple of the meta object's operator parameters (i.e. "rands")."""
         if getattr(self, "_rands", None) is not None:
             return self._rands
 
-        self._rands = tuple(getattr(self, s) for s in getattr(self, "__all_props__", ()))
+        all_props = getattr(self, "__all_props__", None)
+
+        if all_props:
+            self._rands = tuple(getattr(self, s) for s in all_props)
+        else:
+            raise NotImplementedError()
 
         return self._rands
 
@@ -282,7 +289,7 @@ def reify(self):
         if self.obj is not None and not isinstance(self.obj, Var):
             return self.obj
         else:
-            reified_rands, any_unreified = meta_reify_iter(self.rands())
+            reified_rands, any_unreified = meta_reify_iter(self.rands)
 
             # If not all the rands reified, then create another meta
             # object--albeit one with potentially more non-`None` `obj` fields.
@@ -304,8 +311,13 @@ def __eq__(self, other):
 
         assert self.base == other.base
 
-        if self.rands():
-            return all(s == o for s, o in zip(self.rands(), other.rands()))
+        try:
+            rands = self.rands
+        except NotImplementedError:
+            rands = None
+
+        if rands:
+            return all(s == o for s, o in zip(self.rands, other.rands))
         else:
             return NotImplemented
 
@@ -315,16 +327,19 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
     def __hash__(self):
-        return hash((self.base, self.rands()))
+        return hash((self.base, self.rands))
 
     def __str__(self):
         return self.__repr__(show_obj=False, _repr=str)
 
     def __repr__(self, show_obj=True, _repr=meta_repr.repr):
-        rands = self.rands()
+        try:
+            rands = self.rands
+        except NotImplementedError:
+            rands = None
 
         if rands:
-            args = _repr(self.rands())[1:-1]
+            args = _repr(self.rands)[1:-1]
         else:
             args = ""
 
@@ -344,8 +359,13 @@ def _repr_pretty_(self, p, cycle):
             with p.group(2, f"{self.__class__.__name__}(", ")"):
                 p.breakable(sep="")
                 idx = None
-                if hasattr(self, "__all_props__"):
-                    for idx, (name, item) in enumerate(zip(self.__all_props__, self.rands())):
+                try:
+                    rands = self.rands
+                except NotImplementedError:
+                    rands = None
+
+                if rands:
+                    for idx, (name, item) in enumerate(zip(self.__all_props__, rands)):
                         if idx:
                             p.text(",")
                             p.breakable()
@@ -386,7 +406,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @abc.abstractmethod
-    def out_meta_types(self, inputs=None):
+    def output_meta_types(self, inputs=None):
         """Return the types of meta variables this `Op` is expected to produce given the inputs."""
         raise NotImplementedError()
 
@@ -411,18 +431,22 @@ class MetaVariable(MetaSymbol):
 
     @property
     @abc.abstractmethod
-    def operator(self):
-        """Return a meta object representing an operator, if any, capable of producing this variable.
+    def base_operator(self):
+        """Return a meta object representing a base-level operator.
 
         It should be callable with all inputs necessary to reproduce this
-        tensor given by `MetaVariable.inputs`.
+        tensor given by `self.base_arguments`.
         """
         raise NotImplementedError()
 
     @property
     @abc.abstractmethod
-    def inputs(self):
-        """Return the inputs necessary for `MetaVariable.operator` to produced this variable, if any."""
+    def base_arguments(self):
+        """Return the base-level arguments.
+
+        These arguments used in conjunction with the callable
+        `self.base_operator` should re-produce this variable.
+        """
         raise NotImplementedError()
 
 
@@ -431,7 +455,12 @@ def _find_meta_type(obj_type, meta_abs_type):
     obj_cls = None
     while True:
         try:
-            obj_cls = next(filter(lambda t: issubclass(obj_type, t.base), cls.__subclasses__()))
+            obj_cls = next(
+                filter(
+                    lambda t: isinstance(t.base, type) and issubclass(obj_type, t.base),
+                    cls.__subclasses__(),
+                )
+            )
         except StopIteration:
             # The current class is the best fit.
             if cls.base == obj_type:

symbolic_pymc/relations/__init__.py

@@ -7,7 +7,7 @@
 from kanren.term import term, operator, arguments
 from kanren.goals import conso
 
-from ..etuple import etuplize, ExpressionTuple
+from ..etuple import etuplize
 
 
 # Hierarchical models that we recognize.

symbolic_pymc/relations/tensorflow/__init__.py

@@ -0,0 +1,16 @@
+from unification import var
+
+from kanren.facts import fact
+from kanren.assoccomm import commutative, associative
+
+from ...tensorflow.meta import mt, TFlowMetaOperator
+
+
+# TODO: We could use `mt.*.op_def.obj.is_commutative` to capture
+# more/all cases.
+fact(commutative, TFlowMetaOperator(mt.AddV2.op_def, var()))
+fact(commutative, TFlowMetaOperator(mt.AddN.op_def, var()))
+fact(commutative, TFlowMetaOperator(mt.Mul.op_def, var()))
+
+fact(associative, TFlowMetaOperator(mt.AddN.op_def, var()))
+fact(associative, TFlowMetaOperator(mt.AddV2.op_def, var()))

symbolic_pymc/relations/theano/__init__.py

@@ -4,13 +4,22 @@
 
 from kanren import eq
 from kanren.core import lall
+from kanren.facts import fact
+from kanren.assoccomm import commutative, associative
+
 
 from .linalg import buildo
 from ..graph import graph_applyo, seq_apply_anyo
 from ...etuple import etuplize, etuple
 from ...theano.meta import mt
 
 
+fact(commutative, mt.add)
+fact(commutative, mt.mul)
+fact(associative, mt.add)
+fact(associative, mt.mul)
+
+
 def tt_graph_applyo(relation, a, b, preprocess_graph=partial(etuplize, shallow=True)):
     """Construct a `graph_applyo` goal that judiciously expands a Theano meta graph.