Skip to content

Commit 8ba6fa2

Browse files
Amit Patankartensorflower-gardener
Amit Patankar
authored andcommitted
Fix heap-buffer-overflow issue with tf.raw_ops.SparseSplit.
PiperOrigin-RevId: 371242872 Change-Id: I482bb3d12602c7c3cc9446f97fb9f584bb98e9a4
1 parent 699bff5 commit 8ba6fa2

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

Diff for: tensorflow/core/util/sparse/sparse_tensor.h

+4
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,10 @@ inline Status SparseTensor::Split(const SparseTensor& input_tensor,
527527
for (int i = 0; i < input_tensor.indices().dim_size(0); ++i) {
528528
const int dim = input_tensor.indices().matrix<int64>()(i, split_dim);
529529
int slice_index = GetSliceIndex(dim, split_size, residual);
530+
if (slice_index >= num_values.size()) {
531+
return errors::InvalidArgument("Slice index ", slice_index,
532+
" is larger than num_split.");
533+
}
530534
num_values[slice_index]++;
531535
}
532536

0 commit comments

Comments
 (0)