26 lines
1.1 KiB
Diff
26 lines
1.1 KiB
Diff
|
From da5ff2daf618591f64b2b62d9d9803951b945e9f Mon Sep 17 00:00:00 2001
|
||
|
|
From: Amit Patankar <amitpatankar@google.com>
|
||
|
|
Date: Wed, 28 Apr 2021 11:24:45 -0700
|
||
|
|
Subject: [PATCH] Fix FPE issue with `tf.raw_ops.DenseCountSparseOutput`.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 370946862
|
||
|
|
Change-Id: I3752584ad04aaecb327ff6793a9640ac56acfe7a
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/count_ops.cc | 3 +++
|
||
|
|
1 file changed, 3 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
|
||
|
|
index eeb0b853c7aad..40aa1fe458c1e 100644
|
||
|
|
--- a/tensorflow/core/kernels/count_ops.cc
|
||
|
|
+++ b/tensorflow/core/kernels/count_ops.cc
|
||
|
|
@@ -122,6 +122,9 @@ class DenseCount : public OpKernel {
|
||
|
|
|
||
|
|
int num_batch_elements = 1;
|
||
|
|
for (int i = 0; i < num_batch_dimensions; ++i) {
|
||
|
|
+ OP_REQUIRES(context, data.shape().dim_size(i) != 0,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "Invalid input: Shapes dimension cannot be 0."));
|
||
|
|
num_batch_elements *= data.shape().dim_size(i);
|
||
|
|
}
|
||
|
|
int num_value_elements = data.shape().num_elements() / num_batch_elements;
|