fix crash in single feature tree boosting

interpretml · Dec 26, 2024 · 50d4254 · 50d4254
1 parent a8e1320
commit 50d4254
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/shared/libebm/PartitionOneDimensionalBoosting.cpp b/shared/libebm/PartitionOneDimensionalBoosting.cpp
@@ -293,6 +293,7 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
                            deltaStepMax);
                   }
 
+                  EBM_ASSERT(pUpdateScore < aUpdateScore + cScores * cSlices);
                   *pUpdateScore = static_cast<FloatScore>(updateScore);
                   ++pUpdateScore;
 
@@ -322,6 +323,20 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
             done:;
                EBM_ASSERT(cSamplesTotalDebug == cSamplesExpectedDebug);
 
+               EBM_ASSERT(bNominal || pUpdateScore == aUpdateScore + cScores * cSlices);
+
+               EBM_ASSERT(bNominal || pSplit == cSlices - 1 + pInnerTermUpdate->GetSplitPointer(iDimension));
+
+#ifndef NDEBUG
+               UIntSplit prevDebug = 0;
+               for(size_t iDebug = 0; iDebug < cSlices - 1; ++iDebug) {
+                  UIntSplit curDebug = pInnerTermUpdate->GetSplitPointer(iDimension)[iDebug];
+                  EBM_ASSERT(prevDebug < curDebug);
+                  prevDebug = curDebug;
+               }
+               EBM_ASSERT(prevDebug < cBins);
+#endif
+
                EBM_ASSERT(nullptr == pMissingValueTreeNode || nullptr != pMissingBin);
                if(nullptr != pMissingBin) {
                   EBM_ASSERT(bMissing);
@@ -347,18 +362,6 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
                   } while(pGradientPairEnd != pGradientPair);
                }
 
-               EBM_ASSERT(bNominal || pSplit == cSlices - 1 + pInnerTermUpdate->GetSplitPointer(iDimension));
-
-#ifndef NDEBUG
-               UIntSplit prevDebug = 0;
-               for(size_t iDebug = 0; iDebug < cSlices - 1; ++iDebug) {
-                  UIntSplit curDebug = pInnerTermUpdate->GetSplitPointer(iDimension)[iDebug];
-                  EBM_ASSERT(prevDebug < curDebug);
-                  prevDebug = curDebug;
-               }
-               EBM_ASSERT(prevDebug < cBins);
-#endif
-
                LOG_0(Trace_Verbose, "Exited Flatten");
                return Error_None;
             }
@@ -379,6 +382,7 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
                            if(cBins == iEdge) {
                               // This cut would isolate the missing bin, but missing already has a cut.
                               // We still need to find the missing bin though in the tree, so continue.
+                              bDone = true;
                               break;
                            }
                         }