From 8ba6fa29cd8bf9cef9b718dc31c78c73081f5b31 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Thu, 29 Apr 2021 17:58:08 -0700
Subject: [PATCH] Fix heap-buffer-overflow issue with `tf.raw_ops.SparseSplit`.

PiperOrigin-RevId: 371242872
Change-Id: I482bb3d12602c7c3cc9446f97fb9f584bb98e9a4
---
 tensorflow/core/util/sparse/sparse_tensor.h | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/core/util/sparse/sparse_tensor.h b/tensorflow/core/util/sparse/sparse_tensor.h
index 062226d7699bc..341290dbbc698 100644
--- a/tensorflow/core/util/sparse/sparse_tensor.h
+++ b/tensorflow/core/util/sparse/sparse_tensor.h
@@ -527,6 +527,10 @@ inline Status SparseTensor::Split(const SparseTensor& input_tensor,
   for (int i = 0; i < input_tensor.indices().dim_size(0); ++i) {
     const int dim = input_tensor.indices().matrix<int64>()(i, split_dim);
     int slice_index = GetSliceIndex(dim, split_size, residual);
+    if (slice_index >= num_values.size()) {
+      return errors::InvalidArgument("Slice index ", slice_index,
+                                     " is larger than num_split.");
+    }
     num_values[slice_index]++;
   }