Google Sync: misc bug fixes.

tensorflow_fold/blocks/blocks.py

@@ -761,8 +761,9 @@ def _types_backward(self, child):
     if in_type is None: return
     wires = self._child_input_wire_dict.get(child)
     if wires is None: return
-    if len(wires) == 1: in_type = (in_type,)
-    if len(wires) != len(in_type):
+    if isinstance(in_type, tdt.PyObjectType) or len(wires) == 1:
+      in_type = (in_type,) * len(wires)
+    elif len(wires) != len(in_type):
       raise TypeError('block %s has %d inputs but expects %d' %
                       (child, len(wires), len(in_type)))
     for i, ((block, index), itype) in enumerate(zip(wires, in_type)):
@@ -878,11 +879,11 @@ def connect(self, a, b):
     if not input_wires:
       b.set_input_type(tdt.VoidType())
     elif len(input_wires) > 1:
-      b.set_input_type_class(tdt.TupleType)
+      b.set_input_type_classes(tdt.TupleType, tdt.PyObjectType)
 
     for block, index in input_wires:
       if index is not None:
-        block.set_output_type_class(tdt.TupleType)
+        block.set_output_type_classes(tdt.TupleType)
         if block.output_type is not None:
           arity = len(block.output_type)
           if arity < index:

tensorflow_fold/blocks/blocks_test.py

@@ -109,6 +109,13 @@ def test_from_tensor_raises(self):
         TypeError, 'shape (1, 2, ?) is not fully defined; call set_shape()',
         tdb.FromTensor, tf.placeholder('int32', [1, 2, None]))
 
+  def test_pyobject_composition(self):
+    block = tdb.AllOf(tdb.Identity(), tdb.Identity())
+    self.assertBuildsConst(('foo', 'foo'), block, 'foo')
+    block = (tdb.AllOf(tdb.Identity(), tdb.Identity()) >>
+             (tdb.Scalar(), tdb.Scalar()))
+    self.assertBuilds((42.0, 42.0), block, 42, max_depth=None)
+
   def test_function_composition(self):
     sc = tdb.Scalar()
     fn1 = times_scalar_block(2.0)

tensorflow_fold/blocks/metrics.py

@@ -16,6 +16,7 @@
 from __future__ import division
 from __future__ import print_function
 # import google3
+import six
 import tensorflow_fold.blocks.blocks as tdb
 import tensorflow_fold.blocks.result_types as tdt
 
@@ -79,6 +80,8 @@ class Metric(tdb.Block):
   """
 
   def __init__(self, metric_name):
+    if not isinstance(metric_name, six.string_types):
+      raise TypeError('metric_name must be a string: %s' % (metric_name,))
     self._metric_name = metric_name
     super(Metric, self).__init__(name=str(metric_name),
                                  output_type=tdt.VoidType())

tensorflow_fold/blocks/result_types.py

@@ -392,13 +392,17 @@ def expected(self):
 
   @expected.setter
   def expected(self, type_classes):
-    if not (self.expected is None or all(issubclass(type_class, self._expected)
-                                         for type_class in type_classes)):
+    """Updates expected type classes to their intersection with type_classes."""
+    if self.expected is None:
+      self._expected = type_classes
+      return
+    expected = tuple(t for t in type_classes if issubclass(t, self._expected))
+    if not expected:
       new_types_str = ' or '.join(x.__name__ for x in type_classes)
       old_types_str = ' or '.join(x.__name__ for x in self._expected)
       raise TypeError('bad %s type %s for %s, expected %s' % (
           self._name, new_types_str, self._contained_by, old_types_str))
-    self._expected = type_classes
+    self._expected = expected
 
 
 class IOBase(object):
@@ -460,21 +464,20 @@ def set_input_type(self, input_type):
     if input_type is not None: self._input_type.value = input_type
     return self
 
-  def set_input_type_class(self, input_type_class):
-    """Updates the type class of the input type.
+  def set_input_type_classes(self, *input_type_classes):
+    """Updates the type classes of the input type.
 
     Args:
-      input_type_class: A type class, or None.
+      *input_type_classes: A tuple of type classes.
 
     Returns:
       `self`
 
     Raises:
-      TypeError: If `input_type_class` is not compatible with the current input
-        type or its expected type classes.
+      TypeError: If `input_type_classes` are not compatible with the current
+        input type or its expected type classes.
     """
-    if input_type_class is not None:
-      self._input_type.expected = (input_type_class,)
+    self._input_type.expected = input_type_classes
     return self
 
   def _update_input_type(self):
@@ -497,21 +500,20 @@ def set_output_type(self, output_type):
     if output_type is not None: self._output_type.value = output_type
     return self
 
-  def set_output_type_class(self, output_type_class):
+  def set_output_type_classes(self, *output_type_classes):
     """Updates the type class of the output type.
 
     Args:
-      output_type_class: A type class, or None.
+      *output_type_classes: A tuple of type classes.
 
     Returns:
       `self`
 
     Raises:
-      TypeError: If `output_type_class` is not compatible with the current
+      TypeError: If `output_type_classes` are not compatible with the current
         output type or its expected type classes.
     """
-    if output_type_class is not None:
-      self._output_type.expected = (output_type_class,)
+    self._output_type.expected = output_type_classes
     return self
 
   def _update_output_type(self):

tensorflow_fold/g3doc/py/td.md

@@ -2707,15 +2707,15 @@ Updates the input type.
 
 - - -
 
-<a name="td.IOBase.set_input_type_class"></a>
-#### `td.IOBase.set_input_type_class(input_type_class)`
+<a name="td.IOBase.set_input_type_classes"></a>
+#### `td.IOBase.set_input_type_classes(*input_type_classes)`
 
-Updates the type class of the input type.
+Updates the type classes of the input type.
 
 ##### Args:
 
 
-*  <b>`input_type_class`</b>: A type class, or None.
+*  <b>`*input_type_classes`</b>: A tuple of type classes.
 
 ##### Returns:
 
@@ -2724,8 +2724,8 @@ Updates the type class of the input type.
 ##### Raises:
 
 
-*  <b>`TypeError`</b>: If `input_type_class` is not compatible with the current input
-    type or its expected type classes.
+*  <b>`TypeError`</b>: If `input_type_classes` are not compatible with the current
+    input type or its expected type classes.
 
 
 - - -
@@ -2775,15 +2775,15 @@ Updates the output type.
 
 - - -
 
-<a name="td.IOBase.set_output_type_class"></a>
-#### `td.IOBase.set_output_type_class(output_type_class)`
+<a name="td.IOBase.set_output_type_classes"></a>
+#### `td.IOBase.set_output_type_classes(*output_type_classes)`
 
 Updates the type class of the output type.
 
 ##### Args:
 
 
-*  <b>`output_type_class`</b>: A type class, or None.
+*  <b>`*output_type_classes`</b>: A tuple of type classes.
 
 ##### Returns:
 
@@ -2792,7 +2792,7 @@ Updates the type class of the output type.
 ##### Raises:
 
 
-*  <b>`TypeError`</b>: If `output_type_class` is not compatible with the current
+*  <b>`TypeError`</b>: If `output_type_classes` are not compatible with the current
     output type or its expected type classes.
 
 

tensorflow_fold/loom/BUILD

@@ -139,4 +139,3 @@ fold_py_test(
         "@org_tensorflow//tensorflow:tensorflow_py",
     ],
 )
-

tensorflow_fold/loom/benchmarks/iclr_2017_benchmark.py

@@ -11,8 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""This is the benchmark code used in our ICLR 2017 paper, entitled
-"Deep Learning with Dynamic Computation graphs."
+"""This is the benchmark code used in the ICLR 2017 paper.
+
+The paper is entitled Deep Learning with Dynamic Computation graphs."
 """
 
 from __future__ import absolute_import
@@ -22,12 +23,15 @@
 import logging
 import random
 import time
+
+# import google3
+
 import numpy as np
 import six
 
 import tensorflow as tf
 
-from tensorflow_fold.loom import loom
+from tensorflow_fold.public import loom
 
 tf.flags.DEFINE_integer("vector_size", 1024, "Size of tree RNN output vector.")
 tf.flags.DEFINE_integer("tree_size", 128, "Size of trees to test.")
@@ -46,8 +50,8 @@
 FLAGS = tf.flags.FLAGS
 
 
-logging.basicConfig(format='%(asctime)s %(message)s')
-_logger = logging.getLogger('benchmark')
+logging.basicConfig(format="%(asctime)s %(message)s")
+_logger = logging.getLogger("benchmark")
 _logger.setLevel(logging.INFO)
 
 
@@ -58,18 +62,21 @@ def make_random_tree(size):
   r = random.randint(1, size-1)
   return (make_random_tree(r), make_random_tree(size-r))
 
+
 def make_sequence_tree(size):
   """Make a maximally unbalanced tree (a sequence) with size nodes."""
   if size <= 1:
     return 0
   return (make_sequence_tree(size-1), 0)
 
+
 def make_balanced_tree(size):
   """Make a balanced binary tree with size nodes, where size is a power of 2."""
   if size <= 1:
     return 0
   return (make_balanced_tree(size/2), make_balanced_tree(size/2))
 
+
 def make_input_tree(size):
   """Make a tree based on the value of the tree_type flag."""
   if FLAGS.tree_type == "sequence":
@@ -82,10 +89,11 @@ def make_input_tree(size):
 
 
 def index_type():
-  return loom.TypeShape('int32', ())
+  return loom.TypeShape("int32", ())
+
 
 def vector_type():
-  return loom.TypeShape('float32', (FLAGS.vector_size,))
+  return loom.TypeShape("float32", (FLAGS.vector_size,))
 
 
 class LeafOp(loom.LoomOp):
@@ -128,8 +136,8 @@ def tree_fc(self, left, right):
     # A simple tree RNN with a single fully connected layer.
     if self._weights is None:
       with tf.variable_scope(self._vscope):
-        self._weights = tf.get_variable("weights",
-            [FLAGS.vector_size*2, FLAGS.vector_size],
+        self._weights = tf.get_variable(
+            "weights", [FLAGS.vector_size*2, FLAGS.vector_size],
             initializer=tf.uniform_unit_scaling_initializer(1.43))
         self._bias = tf.get_variable("bias", [FLAGS.vector_size],
                                      initializer=tf.zeros_initializer())
@@ -141,13 +149,13 @@ def tree_lstm(self, left, right):
     # A variation on the tree LSTM -- we add an extra hidden layer.
     if self._weights is None:
       with tf.variable_scope(self._vscope):
-        self._weights_0 = tf.get_variable("weights_0",
-            [FLAGS.vector_size*2, FLAGS.vector_size],
+        self._weights_0 = tf.get_variable(
+            "weights_0", [FLAGS.vector_size*2, FLAGS.vector_size],
             initializer=tf.uniform_unit_scaling_initializer(1.43))
         self._bias_0 = tf.get_variable("bias_0", [FLAGS.vector_size],
                                        initializer=tf.zeros_initializer())
-        self._weights = tf.get_variable("weights",
-            [FLAGS.vector_size, FLAGS.vector_size*4],
+        self._weights = tf.get_variable(
+            "weights", [FLAGS.vector_size, FLAGS.vector_size*4],
             initializer=tf.uniform_unit_scaling_initializer(1.0))
         self._bias = tf.get_variable("bias", [FLAGS.vector_size*4],
                                      initializer=tf.zeros_initializer())
@@ -166,7 +174,7 @@ def tree_lstm(self, left, right):
 
     ylr = tf.add(tf.multiply(fl, left), tf.multiply(fr, right))
     ygi = tf.multiply(i, g)
-    y   = tf.add(ylr, ygi)
+    y = tf.add(ylr, ygi)
 
     return y
 
@@ -318,7 +326,7 @@ def rand_indices():
       return np.array(
           [self.random_index() for _ in six.moves.xrange(0, self.batch_size)],
           dtype="int32")
-    return { p: rand_indices() for p in self._placeholders }
+    return {p: rand_indices() for p in self._placeholders}
 
 
 class LoomModel(ModelBase):
@@ -345,8 +353,8 @@ def build_model(self):
     # Create a dictionary of the LoomOps that the model uses.
     named_tensors = {}
     named_ops = {
-      "leaf": self._leaf_op,
-      "non_terminal": self._non_terminal_op
+        "leaf": self._leaf_op,
+        "non_terminal": self._non_terminal_op
     }
     # Make a random tree.
     self._tree = make_input_tree(FLAGS.tree_size)
@@ -403,7 +411,7 @@ def build_feed_dict_with_serialize_and_merge(self):
       serialized_trees.append(weaver.serialize())
 
     # Pass the serialized trees as the input tensors.
-    return { self._loom.input_tensor: serialized_trees }
+    return {self._loom.input_tensor: serialized_trees}
 
   def traverse_tree(self, node, weaver):
     # Recursive function to invoke a LoomOp on each node in the tree.
@@ -428,7 +436,7 @@ def test_model(model_class, *args):
 
   for batch_size in batch_size_list:
     test_results[batch_size] = ([], [])
-    for epoch in six.moves.xrange(0, FLAGS.num_epochs):
+    for _ in six.moves.xrange(0, FLAGS.num_epochs):
       model = model_class(batch_size, *args)
       model.run()
       test_results[batch_size][0].extend(model.elapsed_times)
@@ -482,7 +490,7 @@ def avg(lst):
     tree_times = [t/b for t in times]
     avg_time = avg(tree_times)
     _logger.info("Batch size: %d | tree time: %f, speedup: %f", b,
-                    avg_time, baseline_tree_time/avg_time)
+                 avg_time, baseline_tree_time/avg_time)
 
 
 def main(unused_argv):
@@ -500,7 +508,7 @@ def main(unused_argv):
 
   _logger.info("====================================================")
   _logger.info("Num epochs: %d; repeats per epoch %d",
-                  FLAGS.num_epochs, FLAGS.num_repeats)
+               FLAGS.num_epochs, FLAGS.num_repeats)
   _logger.info("Model type: %s, %s", model_type, FLAGS.tree_type)
   _logger.info("Vector size: %d", FLAGS.vector_size)
   _logger.info("Tree size: %d", FLAGS.tree_size)

tensorflow_fold/run_all_examples.sh

@@ -7,7 +7,7 @@ set -e
 bazel build "$@" tensorflow_fold/...
 
 # loom benchmark
-./bazel-bin/tensorflow_fold/loom/loom_benchmark
+./bazel-bin/tensorflow_fold/loom/benchmarks/iclr_2017_benchmark
 
 # loom calculator
 TMP=$(mktemp -d)