  void Compute(OpKernelContext* const context) override {
    core::RefCountPtr<BoostedTreesEnsembleResource> resource;
    // Get the resource.
    OP_REQUIRES_OK(context, LookupResource(context, HandleFromInput(context, 0),
                                           &resource));

    // Get the inputs.
    OpInputList bucketized_features_list;
    OP_REQUIRES_OK(context, context->input_list("bucketized_features",
                                                &bucketized_features_list));
    std::vector<tensorflow::TTypes<int32>::ConstMatrix> bucketized_features;
    bucketized_features.reserve(bucketized_features_list.size());
    ConvertVectorsToMatrices(bucketized_features_list, bucketized_features);
    const int batch_size = bucketized_features[0].dimension(0);

    const Tensor* cached_tree_ids_t;
    OP_REQUIRES_OK(context,
                   context->input("cached_tree_ids", &cached_tree_ids_t));
    const auto cached_tree_ids = cached_tree_ids_t->vec<int32>();

    const Tensor* cached_node_ids_t;
    OP_REQUIRES_OK(context,
                   context->input("cached_node_ids", &cached_node_ids_t));
    const auto cached_node_ids = cached_node_ids_t->vec<int32>();

    // Allocate outputs.
    Tensor* output_partial_logits_t = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output("partial_logits",
                                            {batch_size, logits_dimension_},
                                            &output_partial_logits_t));
    auto output_partial_logits = output_partial_logits_t->matrix<float>();

    Tensor* output_tree_ids_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output("tree_ids", {batch_size},
                                                     &output_tree_ids_t));
    auto output_tree_ids = output_tree_ids_t->vec<int32>();

    Tensor* output_node_ids_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output("node_ids", {batch_size},
                                                     &output_node_ids_t));
    auto output_node_ids = output_node_ids_t->vec<int32>();

    // Indicate that the latest tree was used.
    const int32 latest_tree = resource->num_trees() - 1;

    if (latest_tree < 0) {
      // Ensemble was empty. Output the very first node.
      output_node_ids.setZero();
      output_tree_ids = cached_tree_ids;
      // All the predictions are zeros.
      output_partial_logits.setZero();
    } else {
      output_tree_ids.setConstant(latest_tree);
      auto do_work = [&resource, &bucketized_features, &cached_tree_ids,
                      &cached_node_ids, &output_partial_logits,
                      &output_node_ids, latest_tree,
                      this](int64 start, int64 end) {
        for (int32 i = start; i < end; ++i) {
          int32 tree_id = cached_tree_ids(i);
          int32 node_id = cached_node_ids(i);
          std::vector<float> partial_tree_logits(logits_dimension_, 0.0);

          if (node_id >= 0) {
            // If the tree was pruned, returns the node id into which the
            // current_node_id was pruned, as well the correction of the cached
            // logit prediction.
            resource->GetPostPruneCorrection(tree_id, node_id, &node_id,
                                             &partial_tree_logits);
            // Logic in the loop adds the cached node value again if it is a
            // leaf. If it is not a leaf anymore we need to subtract the old
            // node's value. The following logic handles both of these cases.
            const auto& node_logits = resource->node_value(tree_id, node_id);
            if (!node_logits.empty()) {
              DCHECK_EQ(node_logits.size(), logits_dimension_);
              for (int32 j = 0; j < logits_dimension_; ++j) {
                partial_tree_logits[j] -= node_logits[j];
              }
            }
          } else {
            // No cache exists, start from the very first node.
            node_id = 0;
          }
          std::vector<float> partial_all_logits(logits_dimension_, 0.0);
          while (true) {
            if (resource->is_leaf(tree_id, node_id)) {
              const auto& leaf_logits = resource->node_value(tree_id, node_id);
              DCHECK_EQ(leaf_logits.size(), logits_dimension_);
              // Tree is done
              const float tree_weight = resource->GetTreeWeight(tree_id);
              for (int32 j = 0; j < logits_dimension_; ++j) {
                partial_all_logits[j] +=
                    tree_weight * (partial_tree_logits[j] + leaf_logits[j]);
                partial_tree_logits[j] = 0;
              }
              // Stop if it was the latest tree.
              if (tree_id == latest_tree) {
                break;
              }
              // Move onto other trees.
              ++tree_id;
              node_id = 0;
            } else {
              node_id =
                  resource->next_node(tree_id, node_id, i, bucketized_features);
            }
          }
          output_node_ids(i) = node_id;
          for (int32 j = 0; j < logits_dimension_; ++j) {
            output_partial_logits(i, j) = partial_all_logits[j];
          }
        }
      };
      // 30 is the magic number. The actual value might be a function of (the
      // number of layers) * (cpu cycles spent on each layer), but this value
      // would work for many cases. May be tuned later.
      const int64 cost = 30;
      thread::ThreadPool* const worker_threads =
          context->device()->tensorflow_cpu_worker_threads()->workers;
      Shard(worker_threads->NumThreads(), worker_threads, batch_size,
            /*cost_per_unit=*/cost, do_work);
    }
  }
  void Compute(OpKernelContext* const context) override {
    core::RefCountPtr<BoostedTreesEnsembleResource> resource;
     
    OP_REQUIRES_OK(context, LookupResource(context, HandleFromInput(context, 0),
                                           &resource));

     
    OpInputList bucketized_features_list;
    OP_REQUIRES_OK(context, context->input_list("bucketized_features",
                                                &bucketized_features_list));
    std::vector<tensorflow::TTypes<int32>::ConstMatrix> bucketized_features;
    bucketized_features.reserve(bucketized_features_list.size());
    ConvertVectorsToMatrices(bucketized_features_list, bucketized_features);
    const int batch_size = bucketized_features[0].dimension(0);

    const Tensor* cached_tree_ids_t;
    OP_REQUIRES_OK(context,
                   context->input("cached_tree_ids", &cached_tree_ids_t));
    const auto cached_tree_ids = cached_tree_ids_t->vec<int32>();

    const Tensor* cached_node_ids_t;
    OP_REQUIRES_OK(context,
                   context->input("cached_node_ids", &cached_node_ids_t));
    const auto cached_node_ids = cached_node_ids_t->vec<int32>();

     
    Tensor* output_partial_logits_t = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output("partial_logits",
                                            {batch_size, logits_dimension_},
                                            &output_partial_logits_t));
    auto output_partial_logits = output_partial_logits_t->matrix<float>();

    Tensor* output_tree_ids_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output("tree_ids", {batch_size},
                                                     &output_tree_ids_t));
    auto output_tree_ids = output_tree_ids_t->vec<int32>();

    Tensor* output_node_ids_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output("node_ids", {batch_size},
                                                     &output_node_ids_t));
    auto output_node_ids = output_node_ids_t->vec<int32>();

     
    const int32 latest_tree = resource->num_trees() - 1;

    if (latest_tree < 0) {
       
      output_node_ids.setZero();
      output_tree_ids = cached_tree_ids;
       
      output_partial_logits.setZero();
    } else {
      output_tree_ids.setConstant(latest_tree);
      auto do_work = [&resource, &bucketized_features, &cached_tree_ids,
                      &cached_node_ids, &output_partial_logits,
                      &output_node_ids, latest_tree,
                      this](int64 start, int64 end) {
        for (int32 i = start; i < end; ++i) {
          int32 tree_id = cached_tree_ids(i);
          int32 node_id = cached_node_ids(i);
          std::vector<float> partial_tree_logits(logits_dimension_, 0.0);

          if (node_id >= 0) {
             
             
             
            resource->GetPostPruneCorrection(tree_id, node_id, &node_id,
                                             &partial_tree_logits);
             
             
             
            const auto& node_logits = resource->node_value(tree_id, node_id);
            if (!node_logits.empty()) {
              DCHECK_EQ(node_logits.size(), logits_dimension_);
              for (int32 j = 0; j < logits_dimension_; ++j) {
                partial_tree_logits[j] -= node_logits[j];
              }
            }
          } else {
             
            node_id = 0;
          }
          std::vector<float> partial_all_logits(logits_dimension_, 0.0);
          while (true) {
            if (resource->is_leaf(tree_id, node_id)) {
              const auto& leaf_logits = resource->node_value(tree_id, node_id);
              DCHECK_EQ(leaf_logits.size(), logits_dimension_);
               
              const float tree_weight = resource->GetTreeWeight(tree_id);
              for (int32 j = 0; j < logits_dimension_; ++j) {
                partial_all_logits[j] +=
                    tree_weight * (partial_tree_logits[j] + leaf_logits[j]);
                partial_tree_logits[j] = 0;
              }
               
              if (tree_id == latest_tree) {
                break;
              }
               
              ++tree_id;
              node_id = 0;
            } else {
              node_id =
                  resource->next_node(tree_id, node_id, i, bucketized_features);
            }
          }
          output_node_ids(i) = node_id;
          for (int32 j = 0; j < logits_dimension_; ++j) {
            output_partial_logits(i, j) = partial_all_logits[j];
          }
        }
      };
       
       
       
      const int64 cost = 30;
      thread::ThreadPool* const worker_threads =
          context->device()->tensorflow_cpu_worker_threads()->workers;
      Shard(worker_threads->NumThreads(), worker_threads, batch_size,
             cost, do_work);
    }
  }