boosted_trees: Removed less used AddTreesToEnsembleOp (only used for tests now in which can be replaced by TreeEnsembleDeserializeOp).

PiperOrigin-RevId: 170247658

10 files changed, 20 insertions, 783 deletions

diff --git a/tensorflow/contrib/boosted_trees/BUILD b/tensorflow/contrib/boosted_trees/BUILD
index 30f12d02f2..726a8f692f 100644
--- a/tensorflow/contrib/boosted_trees/BUILD
+++ b/tensorflow/contrib/boosted_trees/BUILD
@@ -28,7 +28,6 @@ package_group(name = "friends")
 cc_library(
     name = "boosted_trees_kernels",
     deps = [
-        ":ensemble_optimizer_ops_kernels",
         ":model_ops_kernels",
         ":prediction_ops_kernels",
         ":quantile_ops_kernels",
@@ -42,7 +41,6 @@ cc_library(
 cc_library(
     name = "boosted_trees_ops_op_lib",
     deps = [
-        ":ensemble_optimizer_ops_op_lib",
         ":model_ops_op_lib",
         ":prediction_ops_op_lib",
         ":quantile_ops_op_lib",
@@ -128,29 +126,6 @@ py_test(
 # Kernel tests
 
 py_test(
-    name = "ensemble_optimizer_ops_test",
-    size = "small",
-    srcs = ["python/kernel_tests/ensemble_optimizer_ops_test.py"],
-    srcs_version = "PY2AND3",
-    tags = [
-        "nomac",  # b/63258195
-    ],
-    deps = [
-        ":ensemble_optimizer_ops_py",
-        ":model_ops_py",
-        "//tensorflow/contrib/boosted_trees/proto:tree_config_proto_py",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:dtypes",
-        "//tensorflow/python:framework_ops",
-        "//tensorflow/python:framework_test_lib",
-        "//tensorflow/python:platform_test",
-        "//tensorflow/python:resources",
-        "//tensorflow/python:variables",
-        "//third_party/py/numpy",
-    ],
-)
-
-py_test(
     name = "model_ops_test",
     size = "small",
     srcs = ["python/kernel_tests/model_ops_test.py"],
@@ -159,7 +134,6 @@ py_test(
         "nomac",  # b/63258195
     ],
     deps = [
-        ":ensemble_optimizer_ops_py",
         ":model_ops_py",
         ":prediction_ops_py",
         "//tensorflow/contrib/boosted_trees/proto:learner_proto_py",
@@ -304,7 +278,6 @@ py_library(
     name = "boosted_trees_ops_py",
     srcs_version = "PY2AND3",
     deps = [
-        ":ensemble_optimizer_ops_py",
         ":model_ops_py",
         ":prediction_ops_py",
         ":quantile_ops_py",
@@ -361,14 +334,12 @@ tf_kernel_library(
 tf_custom_op_library(
     name = "python/ops/_boosted_trees_ops.so",
     srcs = [
-        "kernels/ensemble_optimizer_ops.cc",
         "kernels/model_ops.cc",
         "kernels/prediction_ops.cc",
         "kernels/quantile_ops.cc",
         "kernels/split_handler_ops.cc",
         "kernels/stats_accumulator_ops.cc",
         "kernels/training_ops.cc",
-        "ops/ensemble_optimizer_ops.cc",
         "ops/model_ops.cc",
         "ops/prediction_ops.cc",
         "ops/quantile_ops.cc",
@@ -585,52 +556,6 @@ tf_kernel_library(
     alwayslink = 1,
 )
 
-# Ensemble optimizer ops
-tf_gen_op_libs(
-    op_lib_names = ["ensemble_optimizer_ops"],
-)
-
-tf_gen_op_wrapper_py(
-    name = "gen_ensemble_optimizer_ops_py",
-    out = "python/ops/gen_ensemble_optimizer_ops.py",
-    deps = [
-        ":ensemble_optimizer_ops_op_lib",
-    ],
-)
-
-tf_custom_op_py_library(
-    name = "ensemble_optimizer_ops_py",
-    srcs = ["python/ops/ensemble_optimizer_ops.py"],
-    kernels = [
-        ":ensemble_optimizer_ops_kernels",
-        ":ensemble_optimizer_ops_op_lib",
-    ],
-    srcs_version = "PY2AND3",
-    deps = [
-        ":boosted_trees_ops_loader",
-        ":gen_ensemble_optimizer_ops_py",
-        "//tensorflow/contrib/util:util_py",
-        "//tensorflow/python:framework_for_generated_wrappers",
-    ],
-)
-
-tf_kernel_library(
-    name = "ensemble_optimizer_ops_kernels",
-    srcs = [
-        "kernels/ensemble_optimizer_ops.cc",
-    ],
-    deps = [
-        "//tensorflow/contrib/boosted_trees/lib:utils",
-        "//tensorflow/contrib/boosted_trees/proto:learner_proto_cc",
-        "//tensorflow/contrib/boosted_trees/proto:tree_config_proto_cc",
-        "//tensorflow/contrib/boosted_trees/resources:decision_tree_ensemble_resource",
-        "//tensorflow/core:framework",
-        "//tensorflow/core:framework_headers_lib",
-        "//third_party/eigen3",
-    ],
-    alwayslink = 1,
-)
-
 # Stats Accumulator ops
 tf_gen_op_libs(
     op_lib_names = ["stats_accumulator_ops"],
diff --git a/tensorflow/contrib/boosted_trees/kernels/ensemble_optimizer_ops.cc b/tensorflow/contrib/boosted_trees/kernels/ensemble_optimizer_ops.cc
deleted file mode 100644
index 5cde229010..0000000000
--- a/tensorflow/contrib/boosted_trees/kernels/ensemble_optimizer_ops.cc
+++ /dev/null
@@ -1,243 +0,0 @@
-// Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// 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.
-// =============================================================================
-#include <string>
-#include <vector>
-
-#include "tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.h"
-#include "tensorflow/contrib/boosted_trees/proto/tree_config.pb.h"
-#include "tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h"
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/framework/resource_mgr.h"
-#include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/framework/tensor_shape.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/core/refcount.h"
-#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/platform/types.h"
-
-namespace tensorflow {
-
-using boosted_trees::models::DecisionTreeEnsembleResource;
-using boosted_trees::trees::DecisionTreeEnsembleConfig;
-using boosted_trees::utils::DropoutUtils;
-using errors::InvalidArgument;
-
-namespace {
-
-// Learning rate epsilon.
-const float kLearningRateEps = 1e-8;
-
-}  // namespace
-
-class AddTreesToEnsembleOp : public OpKernel {
- public:
-  explicit AddTreesToEnsembleOp(OpKernelConstruction* const context)
-      : OpKernel(context) {
-    // Ensure feature importance lhs inputs are references.
-    OP_REQUIRES(
-        context,
-        IsRefType(context->input_type(kFeatureColumnUsageCountsHandleIdx)),
-        errors::InvalidArgument(
-            "Feature usage counts lhs input needs to be a ref type"));
-    OP_REQUIRES(context,
-                IsRefType(context->input_type(kFeatureColumnGainsHandleIdx)),
-                errors::InvalidArgument(
-                    "Feature gains lhs input needs to be a ref type"));
-  }
-
-  void Compute(OpKernelContext* const context) override {
-    DecisionTreeEnsembleResource* decision_tree_ensemble_resource;
-    // Create a reference to the underlying resource using the handle.
-    OP_REQUIRES_OK(
-        context, LookupResource(
-                     context, HandleFromInput(context, kTreeEnsembleHandleIdx),
-                     &decision_tree_ensemble_resource));
-    // Lock the resource since we're mutating it.
-    mutex_lock l(*decision_tree_ensemble_resource->get_mutex());
-    // Remove the reference at the end of this scope.
-    core::ScopedUnref unref_me(decision_tree_ensemble_resource);
-
-    // Read feature importance info.
-    mutex_lock fc_usage_counts_mutex_lock(
-        *context->input_ref_mutex(kFeatureColumnUsageCountsHandleIdx));
-    mutex_lock fc_gains_mutex_lock(
-        *context->input_ref_mutex(kFeatureColumnGainsHandleIdx));
-    Tensor fc_usage_counts_lhs_t =
-        context->mutable_input(kFeatureColumnUsageCountsHandleIdx, true);
-    OP_REQUIRES(context,
-                TensorShapeUtils::IsVector(fc_usage_counts_lhs_t.shape()),
-                InvalidArgument("Feature usage counts should be a vector."));
-    OP_REQUIRES(context, fc_usage_counts_lhs_t.IsInitialized(),
-                errors::FailedPrecondition(
-                    "Attempting to use uninitialized variables: ",
-                    requested_input(kFeatureColumnUsageCountsHandleIdx)));
-
-    Tensor fc_gains_lhs_t =
-        context->mutable_input(kFeatureColumnGainsHandleIdx, true);
-    OP_REQUIRES(context, TensorShapeUtils::IsVector(fc_gains_lhs_t.shape()),
-                InvalidArgument("Feature gains should be a vector."));
-    OP_REQUIRES(context, fc_gains_lhs_t.IsInitialized(),
-                errors::FailedPrecondition(
-                    "Attempting to use uninitialized variables: ",
-                    requested_input(kFeatureColumnGainsHandleIdx)));
-
-    const Tensor fc_usage_counts_rhs_t =
-        context->input(kFeatureColumnUsageCountsToAddIdx);
-    OP_REQUIRES(
-        context,
-        fc_usage_counts_lhs_t.shape().IsSameSize(fc_usage_counts_rhs_t.shape()),
-        errors::InvalidArgument(
-            "Shapes of both feature usage counts tensors should match.",
-            " lhs shape= ", fc_usage_counts_lhs_t.shape().DebugString(),
-            " rhs shape= ", fc_usage_counts_rhs_t.shape().DebugString()));
-
-    const Tensor fc_gains_rhs_t = context->input(kFeatureColumnGainsToAddIdx);
-    OP_REQUIRES(context,
-                fc_gains_lhs_t.shape().IsSameSize(fc_gains_rhs_t.shape()),
-                errors::InvalidArgument(
-                    "Shapes of both feature gains tensors should match.",
-                    " lhs shape= ", fc_gains_lhs_t.shape().DebugString(),
-                    " rhs shape= ", fc_gains_rhs_t.shape().DebugString()));
-
-    // Read in info about trees that were dropped.
-    Tensor dropped_trees_info_t = context->input(kDropedTreesInfoTensorIdx);
-    OP_REQUIRES(context,
-                TensorShapeUtils::IsMatrix(dropped_trees_info_t.shape()),
-                InvalidArgument("Dropped trees info should be matrix."));
-
-    const auto& dropout_info = dropped_trees_info_t.matrix<float>();
-
-    // Parse the passed in tree ensemble.
-    Tensor tree_ensemble_config_t = context->input(kEnsembleToAddTensorIdx);
-    OP_REQUIRES(
-        context, TensorShapeUtils::IsScalar(tree_ensemble_config_t.shape()),
-        errors::InvalidArgument("Tree ensemble config must be a scalar."));
-    // Arena increase spatial locality which reduces the average latency to
-    // access memory, as working set of pages will be fewer.
-    // arena has type proto2::Arena*.
-    auto* arena =
-        decision_tree_ensemble_resource->mutable_decision_tree_ensemble()
-            ->GetArena();
-    DecisionTreeEnsembleConfig* ensemble_to_add =
-        protobuf::Arena::CreateMessage<DecisionTreeEnsembleConfig>(arena);
-    OP_REQUIRES(
-        context, ParseProtoUnlimited(ensemble_to_add,
-                                     tree_ensemble_config_t.scalar<string>()()),
-        errors::InvalidArgument("Unable to parse tree ensemble config."));
-
-    auto* mutable_ensemble =
-        decision_tree_ensemble_resource->mutable_decision_tree_ensemble();
-
-    // Read the learning_rate
-    Tensor learning_rate_t = context->input(kLearningRateTensorIdx);
-    OP_REQUIRES(context, TensorShapeUtils::IsScalar(learning_rate_t.shape()),
-                InvalidArgument("Learning rate should be a scalar."));
-
-    const float learning_rate = learning_rate_t.scalar<float>()();
-    if (learning_rate < kLearningRateEps) {
-      return;
-    }
-    // Prepare current weights vec.
-    std::vector<float> current_weights;
-    current_weights.reserve(mutable_ensemble->tree_weights_size());
-    for (const float weight : mutable_ensemble->tree_weights()) {
-      current_weights.push_back(weight);
-    }
-    const int32 num_dropped = dropped_trees_info_t.dim_size(1);
-    std::vector<int> dropped_trees;
-    dropped_trees.reserve(num_dropped);
-    std::vector<float> dropped_trees_original_weights;
-    dropped_trees_original_weights.reserve(num_dropped);
-    for (int i = 0; i < num_dropped; ++i) {
-      dropped_trees.push_back(dropout_info(0, i));
-      dropped_trees_original_weights.push_back(dropout_info(1, i));
-    }
-
-    std::vector<int32> num_updates;
-    num_updates.reserve(mutable_ensemble->tree_metadata_size());
-
-    for (const auto& meta : mutable_ensemble->tree_metadata()) {
-      num_updates.push_back(meta.num_tree_weight_updates());
-    }
-
-    // If there was a dropout, come up with tree weights
-    const bool was_dropout = !dropped_trees.empty();
-    if (was_dropout) {
-      // New tree/s will be added to the end of the ensemble's tree list.
-      const int32 new_tree_index = current_weights.size();
-      DropoutUtils::GetTreesWeightsForAddingTrees(
-          dropped_trees, dropped_trees_original_weights, new_tree_index,
-          ensemble_to_add->trees_size(), &current_weights, &num_updates);
-
-      // Update the weights of trees according to current weights;
-      for (int i = 0; i < mutable_ensemble->trees_size(); ++i) {
-        mutable_ensemble->set_tree_weights(i, current_weights[i]);
-      }
-    }
-
-    // Add the trees from ensemble_to_add to the tree ensemble variable.
-    int i = mutable_ensemble->trees_size();
-    for (auto& tree : *ensemble_to_add->mutable_trees()) {
-      (*mutable_ensemble->add_trees()).Swap(&tree);
-
-      // New trees were updated only once.
-      auto* meta = mutable_ensemble->add_tree_metadata();
-      meta->set_num_tree_weight_updates(1);
-
-      // When we add complete trees to the ensemble in one step, each tree
-      // that's added is final.
-      meta->set_is_finalized(true);
-
-      if (was_dropout) {
-        mutable_ensemble->add_tree_weights(current_weights[i++]);
-      } else {
-        mutable_ensemble->add_tree_weights(learning_rate);
-      }
-    }
-
-    // Update the number of updates.
-    if (was_dropout) {
-      for (int i = 0; i < num_updates.size(); ++i) {
-        mutable_ensemble->mutable_tree_metadata(i)->set_num_tree_weight_updates(
-            num_updates[i]);
-      }
-    }
-
-    // Update feature importance.
-    fc_usage_counts_lhs_t.vec<int64>() += fc_usage_counts_rhs_t.vec<int64>();
-    fc_gains_lhs_t.vec<float>() += learning_rate * fc_gains_rhs_t.vec<float>();
-  }
-
- private:
-  // Input tensor indices.
-  // Note that Op definition changes might cause input indices to need
-  // changing as well.
-  static const int kTreeEnsembleHandleIdx = 0;
-  static const int kEnsembleToAddTensorIdx = 1;
-  static const int kFeatureColumnUsageCountsHandleIdx = 2;
-  static const int kFeatureColumnUsageCountsToAddIdx = 3;
-  static const int kFeatureColumnGainsHandleIdx = 4;
-  static const int kFeatureColumnGainsToAddIdx = 5;
-  static const int kDropedTreesInfoTensorIdx = 6;
-  static const int kLearningRateTensorIdx = 7;
-};
-
-REGISTER_KERNEL_BUILDER(Name("AddTreesToEnsemble").Device(DEVICE_CPU),
-                        AddTreesToEnsembleOp);
-
-}  // namespace tensorflow
diff --git a/tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc b/tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc
deleted file mode 100644
index b5ea5e7849..0000000000
--- a/tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc
+++ /dev/null
@@ -1,44 +0,0 @@
-// Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// 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.
-// =============================================================================
-#include "tensorflow/core/framework/common_shape_fns.h"
-#include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/framework/shape_inference.h"
-
-namespace tensorflow {
-
-REGISTER_OP("AddTreesToEnsemble")
-    .Input("tree_ensemble_handle: resource")
-    .Input("ensemble_to_add: string")
-    .Input("feature_column_usage_counts_handle: Ref(int64)")
-    .Input("feature_column_usage_counts_to_add: int64")
-    .Input("feature_column_gains_handle: Ref(float)")
-    .Input("feature_column_gains_to_add: float")
-    .Input("drop_out_tree_indices_weights: float")
-    .Input("learning_rate: float")
-    .SetShapeFn(shape_inference::NoOutputs)
-    .Doc(R"doc(
-Synchronously adds a tree ensemble to a an existing tree ensemble variable.
-tree_ensemble_handle: Handle to the ensemble variable.
-ensemble_to_add: Serialized DecisionTreeConfig proto of the tree.
-feature_column_usage_counts_handle: Handle to the feature column usage counts variable.
-feature_column_usage_counts_to_add: Rank 1 Tensor holding feature column usage counts to add.
-feature_column_gains_handle: Handle to the feature column gains variable.
-feature_column_gains_to_add: Rank 1 Tensor holding feature column gains to add.
-drop_out_tree_indices_weights: Rank 2 Tensor containing dropped trees indices
-and original weights of those trees during prediction.
-learning_rate: The learning rate that the tuner found for this iteration.
-)doc");
-
-}  // namespace tensorflow
diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/ensemble_optimizer_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/ensemble_optimizer_ops_test.py
deleted file mode 100644
index 842e0caeca..0000000000
--- a/tensorflow/contrib/boosted_trees/python/kernel_tests/ensemble_optimizer_ops_test.py
+++ /dev/null
@@ -1,351 +0,0 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-# ==============================================================================
-"""Tests for the GTFlow ensemble optimization ops.
-
-The tests cover:
-- Adding a newly built tree to an existing ensemble
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-from tensorflow.contrib.boosted_trees.proto import tree_config_pb2
-from tensorflow.contrib.boosted_trees.python.ops import ensemble_optimizer_ops
-from tensorflow.contrib.boosted_trees.python.ops import model_ops
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import resources
-from tensorflow.python.ops import variables
-from tensorflow.python.platform import googletest
-
-
-def _append_to_leaf(leaf, class_id, weight):
-  """Helper method for building tree leaves.
-
-  Appends weight contributions for the given class index to a leaf node.
-
-  Args:
-    leaf: leaf node to append to, int
-    class_id: class Id for the weight update, int
-    weight: weight contribution value, float
-  """
-  leaf.sparse_vector.index.append(class_id)
-  leaf.sparse_vector.value.append(weight)
-
-
-class EnsembleOptimizerOpsTest(test_util.TensorFlowTestCase):
-
-  def setUp(self):
-    """Create an ensemble of 2 trees."""
-    super(EnsembleOptimizerOpsTest, self).setUp()
-    self._tree_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-    # First tree.
-    tree_1 = self._tree_ensemble.trees.add()
-    _append_to_leaf(tree_1.nodes.add().leaf, 0, 0.4)
-    _append_to_leaf(tree_1.nodes.add().leaf, 1, 0.6)
-    # Second tree.
-    tree_2 = self._tree_ensemble.trees.add()
-    _append_to_leaf(tree_2.nodes.add().leaf, 0, 1)
-    _append_to_leaf(tree_2.nodes.add().leaf, 1, 0)
-
-    self._tree_ensemble.tree_weights.append(1.0)
-    self._tree_ensemble.tree_weights.append(1.0)
-
-    meta_1 = self._tree_ensemble.tree_metadata.add()
-    meta_1.num_tree_weight_updates = 2
-    meta_2 = self._tree_ensemble.tree_metadata.add()
-    meta_2.num_tree_weight_updates = 3
-
-    # Ensemble to be added.
-    self._ensemble_to_add = tree_config_pb2.DecisionTreeEnsembleConfig()
-
-    self._tree_to_add = self._ensemble_to_add.trees.add()
-    _append_to_leaf(self._tree_to_add.nodes.add().leaf, 0, 0.3)
-    _append_to_leaf(self._tree_to_add.nodes.add().leaf, 1, 0.7)
-
-  def testWithEmptyEnsemble(self):
-    with self.test_session():
-      # Create an empty ensemble.
-      tree_ensemble_handle = model_ops.tree_ensemble_variable(
-          stamp_token=0, tree_ensemble_config="", name="empty")
-
-      # Create zero feature importance.
-      feature_usage_counts = variables.Variable(
-          initial_value=array_ops.zeros([1], dtypes.int64),
-          name="feature_usage_counts",
-          trainable=False)
-      feature_gains = variables.Variable(
-          initial_value=array_ops.zeros([1], dtypes.float32),
-          name="feature_gains",
-          trainable=False)
-
-      resources.initialize_resources(resources.shared_resources()).run()
-      variables.initialize_all_variables().run()
-
-      with ops.control_dependencies([
-          ensemble_optimizer_ops.add_trees_to_ensemble(
-              tree_ensemble_handle,
-              self._ensemble_to_add.SerializeToString(),
-              feature_usage_counts, [2],
-              feature_gains, [0.4], [[]],
-              learning_rate=1.0)
-      ]):
-        result = model_ops.tree_ensemble_serialize(tree_ensemble_handle)[1]
-
-      # Output.
-      output_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      output_ensemble.ParseFromString(result.eval())
-      self.assertProtoEquals(self._tree_to_add, output_ensemble.trees[0])
-      self.assertEqual(1, len(output_ensemble.trees))
-
-      self.assertAllEqual([1.0], output_ensemble.tree_weights)
-
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[0].num_tree_weight_updates)
-
-      self.assertAllEqual([2], feature_usage_counts.eval())
-      self.assertArrayNear([0.4], feature_gains.eval(), 1e-6)
-
-  def testWithExistingEnsemble(self):
-    with self.test_session():
-      # Create existing tree ensemble.
-      tree_ensemble_handle = model_ops.tree_ensemble_variable(
-          stamp_token=0,
-          tree_ensemble_config=self._tree_ensemble.SerializeToString(),
-          name="existing")
-      # Create non-zero feature importance.
-      feature_usage_counts = variables.Variable(
-          initial_value=np.array([0, 4, 1], np.int64),
-          name="feature_usage_counts",
-          trainable=False)
-      feature_gains = variables.Variable(
-          initial_value=np.array([0.0, 0.3, 0.05], np.float32),
-          name="feature_gains",
-          trainable=False)
-
-      resources.initialize_resources(resources.shared_resources()).run()
-      variables.initialize_all_variables().run()
-      output_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      with ops.control_dependencies([
-          ensemble_optimizer_ops.add_trees_to_ensemble(
-              tree_ensemble_handle,
-              self._ensemble_to_add.SerializeToString(),
-              feature_usage_counts, [1, 2, 0],
-              feature_gains, [0.02, 0.1, 0.0], [[], []],
-              learning_rate=1)
-      ]):
-        output_ensemble.ParseFromString(
-            model_ops.tree_ensemble_serialize(tree_ensemble_handle)[1].eval())
-
-      # Output.
-      self.assertEqual(3, len(output_ensemble.trees))
-      self.assertProtoEquals(self._tree_to_add, output_ensemble.trees[2])
-
-      self.assertAllEqual([1.0, 1.0, 1.0], output_ensemble.tree_weights)
-
-      self.assertEqual(2,
-                       output_ensemble.tree_metadata[0].num_tree_weight_updates)
-      self.assertEqual(3,
-                       output_ensemble.tree_metadata[1].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[2].num_tree_weight_updates)
-      self.assertAllEqual([1, 6, 1], feature_usage_counts.eval())
-      self.assertArrayNear([0.02, 0.4, 0.05], feature_gains.eval(), 1e-6)
-
-  def testWithExistingEnsembleAndDropout(self):
-    with self.test_session():
-      tree_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      # Add 10 trees with some weights.
-      for i in range(0, 10):
-        tree = tree_ensemble.trees.add()
-        _append_to_leaf(tree.nodes.add().leaf, 0, -0.4)
-        tree_ensemble.tree_weights.append(i + 1)
-        meta = tree_ensemble.tree_metadata.add()
-        meta.num_tree_weight_updates = 1
-      tree_ensemble_handle = model_ops.tree_ensemble_variable(
-          stamp_token=0,
-          tree_ensemble_config=tree_ensemble.SerializeToString(),
-          name="existing")
-      # Create non-zero feature importance.
-      feature_usage_counts = variables.Variable(
-          initial_value=np.array([2, 3], np.int64),
-          name="feature_usage_counts",
-          trainable=False)
-      feature_gains = variables.Variable(
-          initial_value=np.array([0.0, 0.3], np.float32),
-          name="feature_gains",
-          trainable=False)
-
-      resources.initialize_resources(resources.shared_resources()).run()
-      variables.initialize_all_variables().run()
-
-      dropped = [1, 6, 8]
-      dropped_original_weights = [2.0, 7.0, 9.0]
-
-      output_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      with ops.control_dependencies([
-          ensemble_optimizer_ops.add_trees_to_ensemble(
-              tree_ensemble_handle,
-              self._ensemble_to_add.SerializeToString(),
-              feature_usage_counts, [1, 2],
-              feature_gains, [0.5, 0.3], [dropped, dropped_original_weights],
-              learning_rate=0.1)
-      ]):
-        output_ensemble.ParseFromString(
-            model_ops.tree_ensemble_serialize(tree_ensemble_handle)[1].eval())
-
-      # Output.
-      self.assertEqual(11, len(output_ensemble.trees))
-      self.assertProtoEquals(self._tree_to_add, output_ensemble.trees[10])
-      self.assertAllClose(4.5, output_ensemble.tree_weights[10])
-
-      self.assertAllClose([1., 1.5, 3., 4., 5., 6., 5.25, 8., 6.75, 10., 4.5],
-                          output_ensemble.tree_weights)
-
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[0].num_tree_weight_updates)
-      self.assertEqual(2,
-                       output_ensemble.tree_metadata[1].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[2].num_tree_weight_updates)
-
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[3].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[4].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[5].num_tree_weight_updates)
-      self.assertEqual(2,
-                       output_ensemble.tree_metadata[6].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[7].num_tree_weight_updates)
-      self.assertEqual(2,
-                       output_ensemble.tree_metadata[8].num_tree_weight_updates)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[9].num_tree_weight_updates)
-      self.assertEqual(
-          1, output_ensemble.tree_metadata[10].num_tree_weight_updates)
-      self.assertAllEqual([3, 5], feature_usage_counts.eval())
-      self.assertArrayNear([0.05, 0.33], feature_gains.eval(), 1e-6)
-
-  def testWithEmptyEnsembleAndShrinkage(self):
-    with self.test_session():
-      # Add shrinkage config.
-      learning_rate = 0.0001
-      tree_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      tree_ensemble_handle = model_ops.tree_ensemble_variable(
-          stamp_token=0,
-          tree_ensemble_config=tree_ensemble.SerializeToString(),
-          name="existing")
-
-      # Create zero feature importance.
-      feature_usage_counts = variables.Variable(
-          initial_value=np.array([0, 0], np.int64),
-          name="feature_usage_counts",
-          trainable=False)
-      feature_gains = variables.Variable(
-          initial_value=np.array([0.0, 0.0], np.float32),
-          name="feature_gains",
-          trainable=False)
-
-      resources.initialize_resources(resources.shared_resources()).run()
-      variables.initialize_all_variables().run()
-
-      output_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      with ops.control_dependencies([
-          ensemble_optimizer_ops.add_trees_to_ensemble(
-              tree_ensemble_handle,
-              self._ensemble_to_add.SerializeToString(),
-              feature_usage_counts, [1, 2],
-              feature_gains, [0.5, 0.3], [[], []],
-              learning_rate=learning_rate)
-      ]):
-        output_ensemble.ParseFromString(
-            model_ops.tree_ensemble_serialize(tree_ensemble_handle)[1].eval())
-
-      # New tree is added with shrinkage weight.
-      self.assertAllClose([learning_rate], output_ensemble.tree_weights)
-      self.assertEqual(1,
-                       output_ensemble.tree_metadata[0].num_tree_weight_updates)
-      self.assertAllEqual([1, 2], feature_usage_counts.eval())
-      self.assertArrayNear([0.5 * learning_rate, 0.3 * learning_rate],
-                           feature_gains.eval(), 1e-6)
-
-  def testWithExistingEnsembleAndShrinkage(self):
-    with self.test_session():
-      # Add shrinkage config.
-      learning_rate = 0.0001
-      tree_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      # Add 10 trees with some weights.
-      for i in range(0, 5):
-        tree = tree_ensemble.trees.add()
-        _append_to_leaf(tree.nodes.add().leaf, 0, -0.4)
-        tree_ensemble.tree_weights.append(i + 1)
-        meta = tree_ensemble.tree_metadata.add()
-        meta.num_tree_weight_updates = 1
-      tree_ensemble_handle = model_ops.tree_ensemble_variable(
-          stamp_token=0,
-          tree_ensemble_config=tree_ensemble.SerializeToString(),
-          name="existing")
-
-      # Create non-zero feature importance.
-      feature_usage_counts = variables.Variable(
-          initial_value=np.array([4, 7], np.int64),
-          name="feature_usage_counts",
-          trainable=False)
-      feature_gains = variables.Variable(
-          initial_value=np.array([0.2, 0.8], np.float32),
-          name="feature_gains",
-          trainable=False)
-
-      resources.initialize_resources(resources.shared_resources()).run()
-      variables.initialize_all_variables().run()
-
-      output_ensemble = tree_config_pb2.DecisionTreeEnsembleConfig()
-      with ops.control_dependencies([
-          ensemble_optimizer_ops.add_trees_to_ensemble(
-              tree_ensemble_handle,
-              self._ensemble_to_add.SerializeToString(),
-              feature_usage_counts, [1, 2],
-              feature_gains, [0.5, 0.3], [[], []],
-              learning_rate=learning_rate)
-      ]):
-        output_ensemble.ParseFromString(
-            model_ops.tree_ensemble_serialize(tree_ensemble_handle)[1].eval())
-
-      # The weights of previous trees stayed the same, new tree (LAST) is added
-      # with shrinkage weight.
-      self.assertAllClose([1.0, 2.0, 3.0, 4.0, 5.0, learning_rate],
-                          output_ensemble.tree_weights)
-
-      # Check that all number of updates are equal to 1 (e,g, no old tree weight
-      # got adjusted.
-      for i in range(0, 6):
-        self.assertEqual(
-            1, output_ensemble.tree_metadata[i].num_tree_weight_updates)
-
-      # Ensure feature importance was aggregated correctly.
-      self.assertAllEqual([5, 9], feature_usage_counts.eval())
-      self.assertArrayNear(
-          [0.2 + 0.5 * learning_rate, 0.8 + 0.3 * learning_rate],
-          feature_gains.eval(), 1e-6)
-
-if __name__ == "__main__":
-  googletest.main()
diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
index 8e62856854..1ee3d71c5a 100644
--- a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
+++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
@@ -30,13 +30,10 @@ import numpy as np
 
 from tensorflow.contrib.boosted_trees.proto import learner_pb2
 from tensorflow.contrib.boosted_trees.proto import tree_config_pb2
-from tensorflow.contrib.boosted_trees.python.ops import ensemble_optimizer_ops
 from tensorflow.contrib.boosted_trees.python.ops import model_ops
 from tensorflow.contrib.boosted_trees.python.ops import prediction_ops
-from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
-from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import resources
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import googletest
@@ -215,51 +212,34 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
     save_path = os.path.join(self.get_temp_dir(), "restore-test")
     with ops.Graph().as_default() as graph:
       with self.test_session(graph) as sess:
-        tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig()
+        # Prepare learner config.
+        learner_config = learner_pb2.LearnerConfig()
+        learner_config.num_classes = 2
 
+        # Add the first tree and save.
+        tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig()
         tree = tree_ensemble_config.trees.add()
         tree_ensemble_config.tree_metadata.add().is_finalized = True
         tree_ensemble_config.tree_weights.append(1.0)
         _append_to_leaf(tree.nodes.add().leaf, 0, -0.1)
-
-        tree_ensemble_config2 = tree_config_pb2.DecisionTreeEnsembleConfig()
-        tree2 = tree_ensemble_config2.trees.add()
-        tree_ensemble_config.tree_weights.append(1.0)
-        _append_to_leaf(tree2.nodes.add().leaf, 0, -1.0)
-
-        tree_ensemble_config3 = tree_config_pb2.DecisionTreeEnsembleConfig()
-        tree3 = tree_ensemble_config3.trees.add()
-        tree_ensemble_config.tree_weights.append(1.0)
-        _append_to_leaf(tree3.nodes.add().leaf, 0, -10.0)
-
-        # Prepare learner config.
-        learner_config = learner_pb2.LearnerConfig()
-        learner_config.num_classes = 2
-
         tree_ensemble_handle = model_ops.tree_ensemble_variable(
             stamp_token=3,
             tree_ensemble_config=tree_ensemble_config.SerializeToString(),
             name="restore_tree")
-        feature_usage_counts = variables.Variable(
-            initial_value=array_ops.zeros([1], dtypes.int64),
-            name="feature_usage_counts",
-            trainable=False)
-        feature_gains = variables.Variable(
-            initial_value=array_ops.zeros([1], dtypes.float32),
-            name="feature_gains",
-            trainable=False)
-
         resources.initialize_resources(resources.shared_resources()).run()
         variables.initialize_all_variables().run()
         my_saver = saver.Saver()
 
+        # Add the second tree and replace the ensemble of the handle.
+        tree2 = tree_ensemble_config.trees.add()
+        tree_ensemble_config.tree_weights.append(1.0)
+        _append_to_leaf(tree2.nodes.add().leaf, 0, -1.0)
+        # Predict to confirm.
         with ops.control_dependencies([
-            ensemble_optimizer_ops.add_trees_to_ensemble(
+            model_ops.tree_ensemble_deserialize(
                 tree_ensemble_handle,
-                tree_ensemble_config2.SerializeToString(),
-                feature_usage_counts, [0],
-                feature_gains, [0], [[]],
-                learning_rate=1)
+                stamp_token=3,
+                tree_ensemble_config=tree_ensemble_config.SerializeToString())
         ]):
           result, _, _ = prediction_ops.gradient_trees_prediction(
               tree_ensemble_handle,
@@ -280,13 +260,15 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
         self.assertEqual(save_path, val)
 
         # Add more trees after saving.
+        tree3 = tree_ensemble_config.trees.add()
+        tree_ensemble_config.tree_weights.append(1.0)
+        _append_to_leaf(tree3.nodes.add().leaf, 0, -10.0)
+        # Predict to confirm.
         with ops.control_dependencies([
-            ensemble_optimizer_ops.add_trees_to_ensemble(
+            model_ops.tree_ensemble_deserialize(
                 tree_ensemble_handle,
-                tree_ensemble_config3.SerializeToString(),
-                feature_usage_counts, [0],
-                feature_gains, [0], [[]],
-                learning_rate=1)
+                stamp_token=3,
+                tree_ensemble_config=tree_ensemble_config.SerializeToString())
         ]):
           result, _, _ = prediction_ops.gradient_trees_prediction(
               tree_ensemble_handle,
diff --git a/tensorflow/contrib/boosted_trees/python/ops/ensemble_optimizer_ops.py b/tensorflow/contrib/boosted_trees/python/ops/ensemble_optimizer_ops.py
deleted file mode 100644
index f7c2e4fe5a..0000000000
--- a/tensorflow/contrib/boosted_trees/python/ops/ensemble_optimizer_ops.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-# ==============================================================================
-"""Split handler custom ops."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# pylint: disable=unused-import
-from tensorflow.contrib.boosted_trees.python.ops import boosted_trees_ops_loader
-# pylint: enable=unused-import
-# pylint: disable=wildcard-import
-from tensorflow.contrib.boosted_trees.python.ops.gen_ensemble_optimizer_ops import *
-# pylint: enable=wildcard-import
diff --git a/tensorflow/contrib/cmake/tf_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake
index bb0d90213a..61c6686ee0 100644
--- a/tensorflow/contrib/cmake/tf_core_kernels.cmake
+++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake
@@ -40,7 +40,6 @@ endif(tensorflow_BUILD_ALL_KERNELS)
 
 if(tensorflow_BUILD_CONTRIB_KERNELS)
   set(tf_contrib_kernels_srcs
-      "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/kernels/ensemble_optimizer_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/kernels/model_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/kernels/prediction_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/kernels/quantile_ops.cc"
@@ -60,7 +59,6 @@ if(tensorflow_BUILD_CONTRIB_KERNELS)
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/lib/learner/stochastic/handlers/sparse-quantized-feature-column-handler.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/lib/models/multiple_additive_trees.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.cc"
-      "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/model_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/prediction_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc"
diff --git a/tensorflow/contrib/cmake/tf_core_ops.cmake b/tensorflow/contrib/cmake/tf_core_ops.cmake
index f27b2aed36..78bccc08a3 100644
--- a/tensorflow/contrib/cmake/tf_core_ops.cmake
+++ b/tensorflow/contrib/cmake/tf_core_ops.cmake
@@ -77,7 +77,6 @@ GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_split_handler "${tensorflow_source_dir
 GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_training "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/training_ops.cc")
 GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_prediction "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/prediction_ops.cc")
 GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_quantiles "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc")
-GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_ensemble_optimzier "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc")
 GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_stats_accumulator "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc")
 GENERATE_CONTRIB_OP_LIBRARY(cudnn_rnn "${tensorflow_source_dir}/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc")
 GENERATE_CONTRIB_OP_LIBRARY(factorization_clustering "${tensorflow_source_dir}/tensorflow/contrib/factorization/ops/clustering_ops.cc")
diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake
index 400f007ee7..441f00e059 100755
--- a/tensorflow/contrib/cmake/tf_python.cmake
+++ b/tensorflow/contrib/cmake/tf_python.cmake
@@ -756,8 +756,6 @@ GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_prediction_ops"
   DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_prediction_ops.py)
 GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_quantiles_ops"
   DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_quantile_ops.py)
-GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_ensemble_optimzier_ops"
-  DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_ensemble_optimizer_ops.py)
 GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_stats_accumulator_ops"
   DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_stats_accumulator_ops.py)
 GENERATE_PYTHON_OP_LIB("contrib_cudnn_rnn_ops"
@@ -1191,4 +1189,3 @@ else()
       WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/tf_python)
   endif(${tensorflow_ENABLE_GPU})
 endif(${tensorflow_TF_NIGHTLY})
-
diff --git a/tensorflow/contrib/makefile/tf_op_files.txt b/tensorflow/contrib/makefile/tf_op_files.txt
index a7f2be9790..ff298e84ad 100644
--- a/tensorflow/contrib/makefile/tf_op_files.txt
+++ b/tensorflow/contrib/makefile/tf_op_files.txt
@@ -1,4 +1,3 @@
-tensorflow/contrib/boosted_trees/ops/ensemble_optimizer_ops.cc
 tensorflow/contrib/boosted_trees/ops/model_ops.cc
 tensorflow/contrib/boosted_trees/ops/prediction_ops.cc
 tensorflow/contrib/boosted_trees/ops/quantile_ops.cc