Fix tensor_forest to properly support float input weights.

Allow tensor_forest inference to pass kwargs through to trees.
Change: 130940805

Author: tensorflower-gardener

tensorflow/contrib/tensor_forest/core/ops/count_extremely_random_stats_op.cc

@@ -527,9 +527,9 @@ class CountExtremelyRandomStats : public OpKernel {
     auto out_leaves = output_leaves->unaligned_flat<int32>();
 
     // <accumulator, class> -> count delta
-    PairMapType<int32> total_delta;
+    PairMapType<float> total_delta;
     // <accumulator, split, class> -> count delta
-    TupleMapType<int32> split_delta;
+    TupleMapType<float> split_delta;
 
     for (int32 i = 0; i < num_data; ++i) {
       out_leaves(i) = results[i].node_indices.back();

tensorflow/contrib/tensor_forest/python/kernel_tests/count_extremely_random_stats_op_test.py

@@ -71,7 +71,7 @@ def testSimple(self):
 
   def testSimpleWeighted(self):
     with self.test_session():
-      input_weights = [1.0, 2.0, 3.0, 4.0]
+      input_weights = [1.5, 2.0, 3.0, 4.0]
       (pcw_node_sums, _, pcw_splits_indices, pcw_splits_sums, _,
        pcw_totals_indices, pcw_totals_sums, _,
        leaves) = (self.ops.count_extremely_random_stats(
@@ -90,12 +90,13 @@ def testSimpleWeighted(self):
            regression=False))
 
       self.assertAllEqual(
-          [[10., 1., 2., 3., 4.], [3., 1., 2., 0., 0.], [7., 0., 0., 3., 4.]],
+          [[10.5, 1.5, 2., 3., 4.],
+           [3.5, 1.5, 2., 0., 0.], [7., 0., 0., 3., 4.]],
           pcw_node_sums.eval())
       self.assertAllEqual([[0, 0, 0], [0, 0, 1]], pcw_splits_indices.eval())
-      self.assertAllEqual([1., 1.], pcw_splits_sums.eval())
+      self.assertAllEqual([1.5, 1.5], pcw_splits_sums.eval())
       self.assertAllEqual([[0, 2], [0, 0], [0, 1]], pcw_totals_indices.eval())
-      self.assertAllEqual([2., 3., 1.], pcw_totals_sums.eval())
+      self.assertAllEqual([2., 3.5, 1.5], pcw_totals_sums.eval())
       self.assertAllEqual([1, 1, 2, 2], leaves.eval())
 
   def testMissingLabel(self):

tensorflow/contrib/tensor_forest/python/tensor_forest.py

@@ -379,13 +379,14 @@ def training_graph(self, input_data, input_labels, data_spec=None,
 
     return control_flow_ops.group(*tree_graphs, name='train')
 
-  def inference_graph(self, input_data, data_spec=None):
+  def inference_graph(self, input_data, data_spec=None, **inference_args):
     """Constructs a TF graph for evaluating a random forest.
 
     Args:
       input_data: A tensor or SparseTensor or placeholder for input data.
       data_spec: A list of tf.dtype values specifying the original types of
         each column.
+      **inference_args: Keyword arguments to pass through to each tree.
 
     Returns:
       The last op in the random forest inference graph.
@@ -397,8 +398,8 @@ def inference_graph(self, input_data, data_spec=None):
         tree_data = input_data
         if self.params.bagged_features:
           tree_data = self._bag_features(i, input_data)
-        probabilities.append(self.trees[i].inference_graph(tree_data,
-                                                           data_spec))
+        probabilities.append(self.trees[i].inference_graph(
+            tree_data, data_spec, **inference_args))
     with ops.device(self.device_assigner.get_device(0)):
       all_predict = array_ops.pack(probabilities)
       return math_ops.div(
@@ -415,7 +416,7 @@ def average_size(self):
     for i in range(self.params.num_trees):
       with ops.device(self.device_assigner.get_device(i)):
         sizes.append(self.trees[i].size())
-    return math_ops.reduce_mean(array_ops.pack(sizes))
+    return math_ops.reduce_mean(math_ops.to_float(array_ops.pack(sizes)))
 
   # pylint: disable=unused-argument
   def training_loss(self, features, labels):