[YDF] Fix computation in LabelHessianNumericalBucket

rstz · copybara-github · commit b14e89dd63c3 · 2026-02-13T07:40:32.000-08:00
PiperOrigin-RevId: 869726553
diff --git a/yggdrasil_decision_forests/learner/decision_tree/splitter_accumulator.h b/yggdrasil_decision_forests/learner/decision_tree/splitter_accumulator.h
@@ -1706,15 +1706,17 @@ struct LabelHessianNumericalBucket {
     if constexpr (weighted) {
       acc->Add(content.sum_gradient, content.sum_hessian, content.sum_weight);
     } else {
-      acc->Add(content.sum_gradient, content.sum_hessian, 1.f);
+      acc->Add(content.sum_gradient, content.sum_hessian,
+               static_cast<float>(count));
     }
   }
 
   void SubToScoreAcc(LabelHessianNumericalScoreAccumulator* acc) const {
     if constexpr (weighted) {
       acc->Sub(content.sum_gradient, content.sum_hessian, content.sum_weight);
     } else {
-      acc->Sub(content.sum_gradient, content.sum_hessian, 1.f);
+      acc->Sub(content.sum_gradient, content.sum_hessian,
+               static_cast<float>(count));
     }
   }
 

Original file line number	Diff line number	Diff line change
`@@ -1706,15 +1706,17 @@ struct LabelHessianNumericalBucket {`
`1706`	`1706`	`if constexpr (weighted) {`
`1707`	`1707`	`acc->Add(content.sum_gradient, content.sum_hessian, content.sum_weight);`
`1708`	`1708`	`} else {`
`1709`		`- acc->Add(content.sum_gradient, content.sum_hessian, 1.f);`
	`1709`	`+ acc->Add(content.sum_gradient, content.sum_hessian,`
	`1710`	`+ static_cast<float>(count));`
`1710`	`1711`	`}`
`1711`	`1712`	`}`
`1712`	`1713`
`1713`	`1714`	`void SubToScoreAcc(LabelHessianNumericalScoreAccumulator* acc) const {`
`1714`	`1715`	`if constexpr (weighted) {`
`1715`	`1716`	`acc->Sub(content.sum_gradient, content.sum_hessian, content.sum_weight);`
`1716`	`1717`	`} else {`
`1717`		`- acc->Sub(content.sum_gradient, content.sum_hessian, 1.f);`
	`1718`	`+ acc->Sub(content.sum_gradient, content.sum_hessian,`
	`1719`	`+ static_cast<float>(count));`
`1718`	`1720`	`}`
`1719`	`1721`	`}`
`1720`	`1722`