From c57c0b9f3a4f8684f3489dd9a9ec627ad8b599f5 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Mon, 19 Apr 2021 11:33:50 -0700
Subject: [PATCH] Fix the segfault in `tf.raw_ops.SparseCountSparseOutput`.

---
 tensorflow/core/kernels/count_ops.cc | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
index b7bb3ed9..67aafebe 100644
--- a/tensorflow/core/kernels/count_ops.cc
+++ b/tensorflow/core/kernels/count_ops.cc
@@ -200,9 +200,17 @@ class SparseCount : public OpKernel {
                     "The shape argument requires at least one element."));
 
     bool is_1d = shape.NumElements() == 1;
-    int num_batches = is_1d ? 1 : shape.flat<int64>()(0);
+    auto shape_vector = shape.flat<int64>();
+    int num_batches = is_1d ? 1 : shape_vector(0);
     int num_values = values.NumElements();
 
+    for (int b = 0; b < shape_vector.size(); b++) {
+      OP_REQUIRES(context, shape_vector(b) >= 0,
+		  errors::InvalidArgument(
+		      "Elements in dense_shape must be >= 0. Instead got:",
+		      shape.DebugString()));
+    }
+
     OP_REQUIRES(context, num_values == indices.shape().dim_size(0),
                 errors::InvalidArgument(
                     "Number of values must match first dimension of indices.",
-- 
2.23.0