void Compute(OpKernelContext* const context) override {
    // node_id_range
    const Tensor* node_id_range_t;
    OP_REQUIRES_OK(context, context->input("node_id_range", &node_id_range_t));
    const auto node_id_range = node_id_range_t->vec<int32>();
    OP_REQUIRES(
        context, node_id_range_t->dims() == 1,
        errors::InvalidArgument("node_id_range must be a rank 1 tensor, but "
                                "given node_id_range has dims of ",
                                node_id_range_t->dims()));
    OP_REQUIRES(context, node_id_range_t->dim_size(0) == 2,
                errors::InvalidArgument(
                    "node_id_range must be a rank 1 tensor with shape=[2], but "
                    "given node_id_range has shape ",
                    node_id_range_t->dim_size(0), " on its first dim"));
    const int32_t node_id_first = node_id_range(0);  // Inclusive.
    const int32_t node_id_last = node_id_range(1);   // Exclusive.

    // Get stats_summaries_list.
    OpInputList stats_summaries_list;
    OP_REQUIRES_OK(context, context->input_list("stats_summaries_list",
                                                &stats_summaries_list));

    // Infer dimensions of a stats_summary.
    DCHECK_GT(stats_summaries_list.size(), 0);
    const int32_t feature_dims = stats_summaries_list[0].dim_size(1);
    // The last bucket is for default/missing value.
    const int32_t num_buckets = stats_summaries_list[0].dim_size(2) - 1;
    const int32_t logits_dim = logits_dim_;
    const int32_t hessian_dim =
        stats_summaries_list[0].dim_size(3) - logits_dim;
    OP_REQUIRES(context, hessian_dim > 0,
                errors::InvalidArgument("hessian dim should be < 0, got ",
                                        hessian_dim));
    OP_REQUIRES(context, hessian_dim <= logits_dim * logits_dim,
                errors::InvalidArgument(
                    "hessian dim should be <= ", logits_dim * logits_dim,
                    " but got: ", hessian_dim));

    // Vector of stats_summaries; each element is stats for feature of shape
    // [max_splits, feature_dim, num_buckets, logits_dim + hessian_dim].
    std::vector<TTypes<float, 4>::ConstTensor> stats_summaries;
    DCHECK_EQ(stats_summaries_list.size(), num_features_);
    stats_summaries.reserve(num_features_);
    for (const auto& tensor : stats_summaries_list) {
      stats_summaries.emplace_back(tensor.tensor<float, 4>());
    }

    // Split types.
    const Tensor* split_types_t;
    OP_REQUIRES_OK(context, context->input("split_types", &split_types_t));
    const auto split_types = split_types_t->vec<tstring>();
    DCHECK_EQ(split_types.size(), num_features_);
    // Validate.
    for (int i = 0; i < num_features_; ++i) {
      if (!(split_types(i) == kInequalitySplit ||
            split_types(i) == kEqualitySplit)) {
        OP_REQUIRES_OK(
            context,
            errors::Aborted(
                "Operation received an exception: Incorrect split type"));
      }
    }
    // Feature ids.
    const Tensor* candidate_feature_ids_t;
    OP_REQUIRES_OK(context, context->input("candidate_feature_ids",
                                           &candidate_feature_ids_t));
    const auto candidate_feature_ids = candidate_feature_ids_t->vec<int32>();
    DCHECK_EQ(candidate_feature_ids.size(), num_features_);

    // L1, L2, tree_complexity, min_node_weight.
    const Tensor* l1_t;
    OP_REQUIRES_OK(context, context->input("l1", &l1_t));
    const auto l1 = l1_t->scalar<float>()();
    DCHECK_GE(l1, 0);
    if (logits_dim_ > 1) {
      // Multi-class L1 regularization not supported yet.
      DCHECK_EQ(l1, 0);
    }
    const Tensor* l2_t;
    OP_REQUIRES_OK(context, context->input("l2", &l2_t));
    const auto l2 = l2_t->scalar<float>()();
    DCHECK_GE(l2, 0);
    const Tensor* tree_complexity_t;
    OP_REQUIRES_OK(context,
                   context->input("tree_complexity", &tree_complexity_t));
    const auto tree_complexity = tree_complexity_t->scalar<float>()();
    const Tensor* min_node_weight_t;
    OP_REQUIRES_OK(context,
                   context->input("min_node_weight", &min_node_weight_t));
    const auto min_node_weight = min_node_weight_t->scalar<float>()();

    std::vector<int32> output_node_ids;
    std::vector<float> output_gains;
    std::vector<int32> output_feature_ids;
    std::vector<int32> output_feature_dimensions;
    std::vector<int32> output_thresholds;
    std::vector<Eigen::VectorXf> output_left_node_contribs;
    std::vector<Eigen::VectorXf> output_right_node_contribs;
    std::vector<string> output_split_types;

    // TODO(tanzheny) parallelize the computation.
    // Iterate each node and find the best gain per node.
    float parent_gain;
    for (int32_t node_id = node_id_first; node_id < node_id_last; ++node_id) {
      float best_gain = std::numeric_limits<float>::lowest();
      int32_t best_bucket;
      int32_t best_f_id;
      int32_t best_f_dim;
      string best_split_type;
      Eigen::VectorXf best_contrib_for_left(logits_dim);
      Eigen::VectorXf best_contrib_for_right(logits_dim);

      // Sum of gradient and hessian. Compute parent gain using first feature.
      ConstMatrixMap stats_mat(&stats_summaries[0](node_id, 0, 0, 0),
                               num_buckets + 1,  // Including default bucket.
                               logits_dim + hessian_dim);
      const Eigen::VectorXf total_grad =
          stats_mat.leftCols(logits_dim).colwise().sum();
      const Eigen::VectorXf total_hess =
          stats_mat.rightCols(hessian_dim).colwise().sum();
      if (total_hess.norm() < min_node_weight) {
        continue;
      }
      Eigen::VectorXf unused(logits_dim);
      CalculateWeightsAndGains(total_grad, total_hess, l1, l2, &unused,
                               &parent_gain);
      for (int f_idx = 0; f_idx < num_features_; ++f_idx) {
        const string split_type = split_types(f_idx);
        TTypes<float, 4>::ConstTensor stats_summary = stats_summaries[f_idx];
        float f_best_gain = std::numeric_limits<float>::lowest();
        int32_t f_best_bucket;
        int32_t f_best_f_dim;
        string f_best_split_type;
        Eigen::VectorXf f_best_contrib_for_left(logits_dim);
        Eigen::VectorXf f_best_contrib_for_right(logits_dim);

        if (split_type == kInequalitySplit) {
          CalculateBestInequalitySplit(
              stats_summary, node_id, feature_dims, logits_dim, hessian_dim,
              num_buckets, min_node_weight, l1, l2, &f_best_gain,
              &f_best_bucket, &f_best_f_dim, &f_best_split_type,
              &f_best_contrib_for_left, &f_best_contrib_for_right);
        } else {
          CalculateBestEqualitySplit(
              stats_summary, total_grad, total_hess, node_id, feature_dims,
              logits_dim, hessian_dim, num_buckets, l1, l2, &f_best_gain,
              &f_best_bucket, &f_best_f_dim, &f_best_split_type,
              &f_best_contrib_for_left, &f_best_contrib_for_right);
        }
        if (f_best_gain > best_gain) {
          best_gain = f_best_gain;
          best_f_id = candidate_feature_ids(f_idx);
          best_f_dim = f_best_f_dim;
          best_split_type = f_best_split_type;
          best_bucket = f_best_bucket;
          best_contrib_for_left = f_best_contrib_for_left;
          best_contrib_for_right = f_best_contrib_for_right;
        }
      }  // For feature id.
      if (best_gain == std::numeric_limits<float>::lowest()) {
        // Do not add the node if no split is found.
        continue;
      }
      output_node_ids.push_back(node_id);
      // Remove the parent gain for the parent node.
      output_gains.push_back(best_gain - parent_gain);
      output_feature_ids.push_back(best_f_id);
      output_feature_dimensions.push_back(best_f_dim);
      // Default direction is fixed for dense splits.
      // TODO(tanzheny) account for default values.
      output_split_types.push_back(best_split_type);
      output_thresholds.push_back(best_bucket);
      output_left_node_contribs.push_back(best_contrib_for_left);
      output_right_node_contribs.push_back(best_contrib_for_right);
    }  // for node id.
    const int num_nodes = output_node_ids.size();
    // output_node_ids
    Tensor* output_node_ids_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output("node_ids", {num_nodes},
                                                     &output_node_ids_t));
    auto output_node_ids_vec = output_node_ids_t->vec<int32>();

    // output_gains
    Tensor* output_gains_t;
    OP_REQUIRES_OK(context, context->allocate_output("gains", {num_nodes},
                                                     &output_gains_t));
    auto output_gains_vec = output_gains_t->vec<float>();

    // output_feature_ids
    Tensor* output_features_ids_t;
    OP_REQUIRES_OK(context, context->allocate_output("feature_ids", {num_nodes},
                                                     &output_features_ids_t));
    auto output_features_vec = output_features_ids_t->vec<int32>();

    // output_feature_dimensions
    Tensor* output_feature_dimension_t;
    OP_REQUIRES_OK(context,
                   context->allocate_output("feature_dimensions", {num_nodes},
                                            &output_feature_dimension_t));
    auto output_feature_dimensions_vec =
        output_feature_dimension_t->vec<int32>();

    // output_thresholds
    Tensor* output_thresholds_t;
    OP_REQUIRES_OK(context, context->allocate_output("thresholds", {num_nodes},
                                                     &output_thresholds_t));
    auto output_thresholds_vec = output_thresholds_t->vec<int32>();

    // output_left_node_contribs
    Tensor* output_left_node_contribs_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "left_node_contribs", {num_nodes, logits_dim},
                                &output_left_node_contribs_t));
    auto output_left_node_contribs_matrix =
        output_left_node_contribs_t->matrix<float>();

    // output_right_node_contribs
    Tensor* output_right_node_contribs_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "right_node_contribs", {num_nodes, logits_dim},
                                &output_right_node_contribs_t));
    auto output_right_node_contribs_matrix =
        output_right_node_contribs_t->matrix<float>();

    // split type
    Tensor* output_split_types_t;
    OP_REQUIRES_OK(
        context, context->allocate_output("split_with_default_directions",
                                          {num_nodes}, &output_split_types_t));
    auto output_split_types_vec = output_split_types_t->vec<tstring>();

    // Sets output tensors from vectors.
    for (int i = 0; i < num_nodes; ++i) {
      output_node_ids_vec(i) = output_node_ids[i];
      output_features_vec(i) = output_feature_ids[i];
      // Adjust the gains to penalize by tree complexity.
      output_gains_vec(i) = output_gains[i] - tree_complexity;
      output_feature_dimensions_vec(i) = output_feature_dimensions[i];
      output_thresholds_vec(i) = output_thresholds[i];
      for (int j = 0; j < logits_dim; ++j) {
        output_left_node_contribs_matrix(i, j) =
            output_left_node_contribs[i][j];
        output_right_node_contribs_matrix(i, j) =
            output_right_node_contribs[i][j];
      }
      output_split_types_vec(i) = output_split_types[i];
    }
  }