  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);

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

    // Return zero logits if it's an empty ensemble.
    if (resource->num_trees() <= 0) {
      output_logits.setZero();
      return;
    }

    const int32 last_tree = resource->num_trees() - 1;
    auto do_work = [&resource, &bucketized_features, &output_logits, last_tree,
                    this](int64 start, int64 end) {
      for (int32 i = start; i < end; ++i) {
        std::vector<float> tree_logits(logits_dimension_, 0.0);
        int32 tree_id = 0;
        int32 node_id = 0;
        while (true) {
          if (resource->is_leaf(tree_id, node_id)) {
            const float tree_weight = resource->GetTreeWeight(tree_id);
            const auto& leaf_logits = resource->node_value(tree_id, node_id);
            DCHECK_EQ(leaf_logits.size(), logits_dimension_);
            for (int32 j = 0; j < logits_dimension_; ++j) {
              tree_logits[j] += tree_weight * leaf_logits[j];
            }
            // Stop if it was the last tree.
            if (tree_id == last_tree) {
              break;
            }
            // Move onto other trees.
            ++tree_id;
            node_id = 0;
          } else {
            node_id =
                resource->next_node(tree_id, node_id, i, bucketized_features);
          }
        }
        for (int32 j = 0; j < logits_dimension_; ++j) {
          output_logits(i, j) = tree_logits[j];
        }
      }
    };
    // 10 is the magic number. The actual number might depend on (the number of
    // layers in the trees) and (cpu cycles spent on each layer), but this
    // value would work for many cases. May be tuned later.
    const int64 cost = (last_tree + 1) * 10;
    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);

     
    Tensor* output_logits_t = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "logits", {batch_size, logits_dimension_},
                                &output_logits_t));
    auto output_logits = output_logits_t->matrix<float>();

     
    if (resource->num_trees() <= 0) {
      output_logits.setZero();
      return;
    }

    const int32 last_tree = resource->num_trees() - 1;
    auto do_work = [&resource, &bucketized_features, &output_logits, last_tree,
                    this](int64 start, int64 end) {
      for (int32 i = start; i < end; ++i) {
        std::vector<float> tree_logits(logits_dimension_, 0.0);
        int32 tree_id = 0;
        int32 node_id = 0;
        while (true) {
          if (resource->is_leaf(tree_id, node_id)) {
            const float tree_weight = resource->GetTreeWeight(tree_id);
            const auto& leaf_logits = resource->node_value(tree_id, node_id);
            DCHECK_EQ(leaf_logits.size(), logits_dimension_);
            for (int32 j = 0; j < logits_dimension_; ++j) {
              tree_logits[j] += tree_weight * leaf_logits[j];
            }
             
            if (tree_id == last_tree) {
              break;
            }
             
            ++tree_id;
            node_id = 0;
          } else {
            node_id =
                resource->next_node(tree_id, node_id, i, bucketized_features);
          }
        }
        for (int32 j = 0; j < logits_dimension_; ++j) {
          output_logits(i, j) = tree_logits[j];
        }
      }
    };
     
     
     
    const int64 cost = (last_tree + 1) * 10;
    thread::ThreadPool* const worker_threads =
        context->device()->tensorflow_cpu_worker_threads()->workers;
    Shard(worker_threads->NumThreads(), worker_threads, batch_size,
           cost, do_work);
  }