Add a Grappler normalize function

Author: brandonwillard

.pylintrc

@@ -263,7 +263,7 @@ ignore-mixin-members=yes
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis. It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=tensorflow.core.framework,tensorflow.python.framework,tensorflow.python.ops.gen_linalg_ops,tensorflow.python.eager.context,tensorflow.compat.v1
+ignored-modules=tensorflow.core.framework,tensorflow.python.framework,tensorflow.python.ops.gen_linalg_ops,tensorflow.python.eager.context,tensorflow.compat.v1,tensorflow.core.protobuf,tensorflow.python.grappler
 
 # List of classes names for which member attributes should not be checked
 # (useful for classes with attributes dynamically set). This supports can work

symbolic_pymc/tensorflow/graph.py

@@ -0,0 +1,56 @@
+import tensorflow as tf
+
+from tensorflow.core.protobuf import config_pb2
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import importer
+from tensorflow.python.framework import meta_graph
+
+from tensorflow.python.grappler import cluster
+from tensorflow.python.grappler import tf_optimizer
+
+
+try:  # pragma: no cover
+    gcluster = cluster.Cluster()
+except tf.errors.UnavailableError:  # pragma: no cover
+    pass
+
+config = config_pb2.ConfigProto()
+
+
+def normalize_tf_graph(graph_output, new_graph=True, verbose=False):
+    """Use grappler to normalize a graph.
+
+    Arguments
+    ---------
+    graph_output: Tensor
+      A tensor we want to consider as "output" of a `FuncGraph`.
+
+    Returns
+    -------
+    The simplified graph.
+    """
+    train_op = graph_output.graph.get_collection_ref(ops.GraphKeys.TRAIN_OP)
+    train_op.clear()
+    train_op.extend([graph_output])
+
+    metagraph = meta_graph.create_meta_graph_def(graph=graph_output.graph)
+
+    optimized_graphdef = tf_optimizer.OptimizeGraph(
+        config, metagraph, verbose=verbose, cluster=gcluster
+    )
+
+    output_name = graph_output.name
+
+    if new_graph:
+        optimized_graph = ops.Graph()
+    else:  # pragma: no cover
+        optimized_graph = ops.get_default_graph()
+        del graph_output
+
+    with optimized_graph.as_default():
+        importer.import_graph_def(optimized_graphdef, name="")
+
+    opt_graph_output = optimized_graph.get_tensor_by_name(output_name)
+
+    return opt_graph_output

symbolic_pymc/tensorflow/meta.py

@@ -6,7 +6,8 @@
 
 from inspect import Parameter, Signature
 
-from collections import OrderedDict, Sequence
+from collections import OrderedDict
+from collections.abc import Sequence
 
 from functools import partial
 

tests/tensorflow/test_graph.py

@@ -0,0 +1,27 @@
+import numpy as np
+import tensorflow as tf
+
+from symbolic_pymc.tensorflow.graph import normalize_tf_graph
+
+from tests.tensorflow import run_in_graph_mode
+
+
+@run_in_graph_mode
+def test_normalize():
+
+    tf.config.optimizer.set_experimental_options(
+        {
+            "shape_optimizations": True,
+            "arithmetic_optimzation": True,
+            "function_optimization": True,
+            "min_graph_nodes": 0,
+        }
+    )
+    with tf.Graph().as_default() as norm_graph:
+        a_tf = tf.compat.v1.placeholder("float")
+        const_log_tf = 0.5 * np.log(2.0 * np.pi) + tf.math.log(a_tf)
+        normal_const_log_tf = normalize_tf_graph(const_log_tf)
+
+        # Grappler appears to put log ops before const
+        assert normal_const_log_tf.op.inputs[0].op.type == "Log"
+        assert normal_const_log_tf.op.inputs[1].op.type == "Const"

tests/tensorflow/test_kanren.py

@@ -1,3 +1,9 @@
+import tensorflow as tf
+
+from tensorflow.python.framework.ops import disable_tensor_equality
+
+from tensorflow_probability import distributions as tfd
+
 from unification import var, unify
 
 from kanren import run, eq, lall
@@ -6,8 +12,12 @@
 
 from symbolic_pymc.meta import enable_lvar_defaults
 from symbolic_pymc.tensorflow.meta import mt
+from symbolic_pymc.tensorflow.graph import normalize_tf_graph
 
 from tests.tensorflow import run_in_graph_mode
+from tests.tensorflow.utils import mt_normal_log_prob
+
+disable_tensor_equality()
 
 
 @run_in_graph_mode
@@ -47,3 +57,36 @@ def test_commutativity():
 
     res = run(0, q, eq_comm(add_1_mt, add_pattern_mt))
     assert res[0] == add_1_mt.base_arguments[0]
+
+
+@run_in_graph_mode
+def test_commutativity_tfp():
+
+    with tf.Graph().as_default():
+        mu_tf = tf.compat.v1.placeholder(tf.float32, name="mu", shape=tf.TensorShape([None]))
+        tau_tf = tf.compat.v1.placeholder(tf.float32, name="tau", shape=tf.TensorShape([None]))
+
+        normal_tfp = tfd.normal.Normal(mu_tf, tau_tf)
+
+        value_tf = tf.compat.v1.placeholder(tf.float32, name="value", shape=tf.TensorShape([None]))
+
+        normal_log_lik = normal_tfp.log_prob(value_tf)
+
+    normal_log_lik_opt = normalize_tf_graph(normal_log_lik)
+
+    with enable_lvar_defaults("names", "node_attrs"):
+        tfp_normal_pattern_mt = mt_normal_log_prob(var(), var(), var())
+
+    normal_log_lik_mt = mt(normal_log_lik)
+    normal_log_lik_opt_mt = mt(normal_log_lik_opt)
+
+    # Our pattern is the form of an unnormalized TFP normal PDF.
+    assert run(0, True, eq(normal_log_lik_mt, tfp_normal_pattern_mt)) == (True,)
+    # Our pattern should *not* match the Grappler-optimized graph, because
+    # Grappler will reorder terms (e.g. the log + constant
+    # variance/normalization term)
+    assert run(0, True, eq(normal_log_lik_opt_mt, tfp_normal_pattern_mt)) == ()
+
+    # XXX: `eq_comm` is, unfortunately, order sensitive!  LHS should be ground.
+    assert run(0, True, eq_comm(normal_log_lik_mt, tfp_normal_pattern_mt)) == (True,)
+    assert run(0, True, eq_comm(normal_log_lik_opt_mt, tfp_normal_pattern_mt)) == (True,)

tests/tensorflow/utils.py

@@ -1,8 +1,12 @@
-from collections.abc import Mapping
+import numpy as np
 
 import tensorflow as tf
 from tensorflow.python.framework import ops
 
+from collections.abc import Mapping
+
+from symbolic_pymc.tensorflow.meta import mt
+
 
 def assert_ops_equal(a, b, compare_fn=lambda a, b: a.op.type == b.op.type):
     if hasattr(a, "op") or hasattr(b, "op"):
@@ -32,3 +36,33 @@ def assert_ops_equal(a, b, compare_fn=lambda a, b: a.op.type == b.op.type):
 
         for i_a, i_b in zip(a_inputs, b_inputs):
             assert_ops_equal(i_a, i_b)
+
+
+def tfp_normal_log_prob(x, loc, scale):
+    """Create a graph of the Grappler-canonicalized form of a TFP normal log-likelihood."""
+    log_unnormalized = -0.5 * tf.math.squared_difference(x / scale, loc / scale)
+    log_normalization = 0.5 * np.log(2.0 * np.pi) + tf.math.log(scale)
+    return log_unnormalized - log_normalization
+
+
+def mt_normal_log_prob(x, loc, scale):
+    """Create a meta graph for Grappler-canonicalized standard or non-standard TFP normal log-likelihoods."""
+    if loc == 0:
+        log_unnormalized_mt = mt(np.array(-0.5, "float32"))
+        log_unnormalized_mt *= mt.squareddifference(
+            mt(np.array(0.0, "float32")),
+            mt.realdiv(x, scale) if scale != 1 else mt.mul(np.array(1.0, "float32"), x),
+        )
+    else:
+        log_unnormalized_mt = mt(np.array(-0.5, "float32"))
+        log_unnormalized_mt *= mt.squareddifference(
+            mt.realdiv(x, scale) if scale != 1 else mt.mul(np.array(1.0, "float32"), x),
+            mt.realdiv(loc, scale) if scale != 1 else mt.mul(np.array(1.0, "float32"), loc),
+        )
+
+    log_normalization_mt = mt((0.5 * np.log(2.0 * np.pi)).astype("float32"))
+
+    if scale != 1:
+        log_normalization_mt = log_normalization_mt + mt.log(scale)
+
+    return log_unnormalized_mt - log_normalization_mt