Fix templating issue

PiperOrigin-RevId: 569470835
google · Sep 29, 2023 · b3b0ca2 · b3b0ca2
1 parent 4e61fbd
commit b3b0ca2
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/yggdrasil_decision_forests/learner/decision_tree/oblique.cc b/yggdrasil_decision_forests/learner/decision_tree/oblique.cc
@@ -374,7 +374,7 @@ absl::StatusOr<SplitSearchResult> EvaluateProjection(
     const InternalTrainConfig& internal_config, const int first_attribute_idx,
     const NodeConstraints& constraints, int8_t monotonic_direction,
     proto::NodeCondition* condition, SplitterPerThreadCache* cache) {
-  const int min_num_obs =
+  const UnsignedExampleIdx min_num_obs =
       dt_config.in_split_min_examples_check() ? dt_config.min_examples() : 1;
 
   // Projection are never missing.
@@ -386,6 +386,7 @@ absl::StatusOr<SplitSearchResult> EvaluateProjection(
 #endif
 
   // Find a good split in the current_projection.
+  // TODO: Why is internal_config not passed along below?
   SplitSearchResult result;
   if constexpr (is_same<LabelStats, ClassificationLabelStats>::value) {
     result = FindSplitLabelClassificationFeatureNumericalCart(

diff --git a/yggdrasil_decision_forests/learner/decision_tree/training.cc b/yggdrasil_decision_forests/learner/decision_tree/training.cc
@@ -2555,6 +2555,30 @@ SplitSearchResult FindSplitLabelHessianRegressionFeatureNumericalCart(
       attribute_idx, condition, &cache->cache_v2);
 }
 
+template SplitSearchResult
+FindSplitLabelHessianRegressionFeatureNumericalCart<true>(
+    const std::vector<UnsignedExampleIdx>& selected_examples,
+    const std::vector<float>& weights, const std::vector<float>& attributes,
+    const std::vector<float>& gradients, const std::vector<float>& hessians,
+    float na_replacement, UnsignedExampleIdx min_num_obs,
+    const proto::DecisionTreeTrainingConfig& dt_config, double sum_gradient,
+    double sum_hessian, double sum_weights, int32_t attribute_idx,
+    const InternalTrainConfig& internal_config,
+    const NodeConstraints& constraints, int8_t monotonic_direction,
+    proto::NodeCondition* condition, SplitterPerThreadCache* cache);
+
+template SplitSearchResult
+FindSplitLabelHessianRegressionFeatureNumericalCart<false>(
+    const std::vector<UnsignedExampleIdx>& selected_examples,
+    const std::vector<float>& weights, const std::vector<float>& attributes,
+    const std::vector<float>& gradients, const std::vector<float>& hessians,
+    float na_replacement, UnsignedExampleIdx min_num_obs,
+    const proto::DecisionTreeTrainingConfig& dt_config, double sum_gradient,
+    double sum_hessian, double sum_weights, int32_t attribute_idx,
+    const InternalTrainConfig& internal_config,
+    const NodeConstraints& constraints, int8_t monotonic_direction,
+    proto::NodeCondition* condition, SplitterPerThreadCache* cache);
+
 template <bool weighted>
 SplitSearchResult
 FindSplitLabelHessianRegressionFeatureDiscretizedNumericalCart(