36 lines
1.3 KiB
Diff
36 lines
1.3 KiB
Diff
|
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
|
||
|
|
|