From c57c0b9f3a4f8684f3489dd9a9ec627ad8b599f5 Mon Sep 17 00:00:00 2001 From: Amit Patankar 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()(0); + auto shape_vector = shape.flat(); + 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