From e84c975313e8e8e38bb2ea118196369c45c51378 Mon Sep 17 00:00:00 2001
From: Laura Pak <lpak@google.com>
Date: Tue, 27 Jul 2021 12:35:03 -0700
Subject: [PATCH] In tf.raw_ops.BoostedTreesSparseCalculateBestFeatureSplit,
 limit stat_dim in stats_summary_indices to under stats_dims in
 stats_summary_shape

PiperOrigin-RevId: 387171191
Change-Id: I83ca8a75b22aa78c037e8b98779da6cced16bfaa
---
 tensorflow/core/kernels/boosted_trees/stats_ops.cc | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorflow/core/kernels/boosted_trees/stats_ops.cc b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
index 851e5b78..032bbfdf 100644
--- a/tensorflow/core/kernels/boosted_trees/stats_ops.cc
+++ b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
@@ -1025,6 +1025,13 @@ class BoostedTreesSparseCalculateBestFeatureSplitOp : public OpKernel {
       const int32 feature_dim = stats_summary_indices(idx, 1);
       const int32 bucket_id = stats_summary_indices(idx, 2);
       const int32 stat_dim = stats_summary_indices(idx, 3);
+      OP_REQUIRES(context, stat_dim < stats_dims,
+		  errors::InvalidArgument(
+		      "Stat dim, the sum of logits dim and hessian dim in "
+		      "stats_summary_indices, cannot be greater than stats "
+		      "dims, the last value in stats_summary_shape, which was ",
+		      stats_dims, ". At index (", idx,
+		      ", 4), stats_summary_indices contains value ", stat_dim));
       std::pair<FeatureMapIterator, bool> const& f_insert_result = f_map.insert(
           FeatureMapIterator::value_type(feature_dim, BucketMap()));
       auto& b_map = f_insert_result.first->second;
-- 
2.27.0