Set up kanren's associative and commutative unify for TensorFlow

Author: brandonwillard

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.
 

symbolic_pymc/tensorflow/meta.py

@@ -210,6 +210,9 @@ def _metatize_tf_eager(obj):
 
     meta._metatize.add((TFlowMetaOpDef.base,), TFlowMetaOpDef._metatize)
 
+    # Apply TF-specific `kanren` settings
+    from ..relations import tensorflow
+
     return meta._metatize
 
 

symbolic_pymc/theano/meta.py

@@ -9,9 +9,6 @@
 
 from unification import var, isvar, Var
 
-from kanren.facts import fact
-from kanren.assoccomm import commutative, associative
-
 from .ops import RandomVariable
 from ..meta import (
     MetaSymbol,
@@ -51,6 +48,9 @@ def load_dispatcher():
     for new_cls in TheanoMetaSymbol.base_subclasses():
         meta._metatize.add((new_cls.base,), new_cls._metatize)
 
+    # Apply TF-specific `kanren` settings
+    from ..relations import theano
+
     return meta._metatize
 
 
@@ -641,8 +641,3 @@ def mt_diag(v, k=0):
 
 
 mt.diag = mt_diag
-
-fact(commutative, mt.add)
-fact(commutative, mt.mul)
-fact(associative, mt.add)
-fact(associative, mt.mul)

tests/tensorflow/test_kanren.py

@@ -0,0 +1,28 @@
+from unification import var
+
+from kanren import run
+from kanren.assoccomm import eq_comm, commutative
+
+from symbolic_pymc.meta import enable_lvar_defaults
+from symbolic_pymc.tensorflow.meta import mt
+
+from tests.tensorflow import run_in_graph_mode
+
+
+@run_in_graph_mode
+def test_commutativity():
+    with enable_lvar_defaults('names'):
+        add_1_mt = mt(1) + mt(2)
+        add_2_mt = mt(2) + mt(1)
+
+    res = run(0, var('q'), commutative(add_1_mt.base_operator))
+    assert res is not False
+
+    res = run(0, var('q'), eq_comm(add_1_mt, add_2_mt))
+    assert res is not False
+
+    with enable_lvar_defaults('names'):
+        add_pattern_mt = mt(2) + var('q')
+
+    res = run(0, var('q'), eq_comm(add_1_mt, add_pattern_mt))
+    assert res[0] == add_1_mt.base_arguments[0]