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

    // We need to get the feature ids used for splitting and the logits after
    // each split. We will use these to calculate the changes in the prediction
    // (contributions) for an arbitrary activation function (done in Python) and
    // attribute them to the associated feature ids. We will store these in
    // a proto below.
    Tensor* output_debug_info_t = nullptr;
    OP_REQUIRES_OK(
        context, context->allocate_output("examples_debug_outputs_serialized",
                                          {batch_size}, &output_debug_info_t));
    // Will contain serialized protos, per example.
    auto output_debug_info = output_debug_info_t->flat<tstring>();
    const int32 last_tree = resource->num_trees() - 1;

    // For each given example, traverse through all trees keeping track of the
    // features used to split and the associated logits at each point along the
    // path. Note: feature_ids has one less value than logits_path because the
    // first value of each logit path will be the bias.
    auto do_work = [&resource, &bucketized_features, &output_debug_info,
                    last_tree](int64 start, int64 end) {
      for (int32 i = start; i < end; ++i) {
        // Proto to store debug outputs, per example.
        boosted_trees::DebugOutput example_debug_info;
        // Initial bias prediction. E.g., prediction based off training mean.
        const auto& tree_logits = resource->node_value(0, 0);
        DCHECK_EQ(tree_logits.size(), 1);
        float tree_logit = resource->GetTreeWeight(0) * tree_logits[0];
        example_debug_info.add_logits_path(tree_logit);
        int32 node_id = 0;
        int32 tree_id = 0;
        int32 feature_id;
        float past_trees_logit = 0;  // Sum of leaf logits from prior trees.
        // Go through each tree and populate proto.
        while (tree_id <= last_tree) {
          if (resource->is_leaf(tree_id, node_id)) {  // Move onto other trees.
            // Accumulate tree_logits only if the leaf is non-root, but do so
            // for bias tree.
            if (tree_id == 0 || node_id > 0) {
              past_trees_logit += tree_logit;
            }
            ++tree_id;
            node_id = 0;
          } else {  // Add to proto.
            // Feature id used to split.
            feature_id = resource->feature_id(tree_id, node_id);
            example_debug_info.add_feature_ids(feature_id);
            // Get logit after split.
            node_id =
                resource->next_node(tree_id, node_id, i, bucketized_features);
            const auto& tree_logits = resource->node_value(tree_id, node_id);
            DCHECK_EQ(tree_logits.size(), 1);
            tree_logit = resource->GetTreeWeight(tree_id) * tree_logits[0];
            // Output logit incorporates sum of leaf logits from prior trees.
            example_debug_info.add_logits_path(tree_logit + past_trees_logit);
          }
        }
        // Set output as serialized proto containing debug info.
        string serialized = example_debug_info.SerializeAsString();
        output_debug_info(i) = serialized;
      }
    };

    // 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_debug_info_t = nullptr;
    OP_REQUIRES_OK(
        context, context->allocate_output("examples_debug_outputs_serialized",
                                          {batch_size}, &output_debug_info_t));
     
    auto output_debug_info = output_debug_info_t->flat<tstring>();
    const int32 last_tree = resource->num_trees() - 1;

     
     
     
     
    auto do_work = [&resource, &bucketized_features, &output_debug_info,
                    last_tree](int64 start, int64 end) {
      for (int32 i = start; i < end; ++i) {
         
        boosted_trees::DebugOutput example_debug_info;
         
        const auto& tree_logits = resource->node_value(0, 0);
        DCHECK_EQ(tree_logits.size(), 1);
        float tree_logit = resource->GetTreeWeight(0) * tree_logits[0];
        example_debug_info.add_logits_path(tree_logit);
        int32 node_id = 0;
        int32 tree_id = 0;
        int32 feature_id;
        float past_trees_logit = 0;   
         
        while (tree_id <= last_tree) {
          if (resource->is_leaf(tree_id, node_id)) {   
             
             
            if (tree_id == 0 || node_id > 0) {
              past_trees_logit += tree_logit;
            }
            ++tree_id;
            node_id = 0;
          } else {   
             
            feature_id = resource->feature_id(tree_id, node_id);
            example_debug_info.add_feature_ids(feature_id);
             
            node_id =
                resource->next_node(tree_id, node_id, i, bucketized_features);
            const auto& tree_logits = resource->node_value(tree_id, node_id);
            DCHECK_EQ(tree_logits.size(), 1);
            tree_logit = resource->GetTreeWeight(tree_id) * tree_logits[0];
             
            example_debug_info.add_logits_path(tree_logit + past_trees_logit);
          }
        }
         
        string serialized = example_debug_info.SerializeAsString();
        output_debug_info(i) = serialized;
      }
    };

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