From f4c364a5d6880557f6f5b6eb5cee2c407f0186b3 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 4 May 2021 18:06:03 -0700 Subject: [PATCH] Fix multiple issues in EditDistance PiperOrigin-RevId: 372033948 Change-Id: Ieb957c29894af05bdfeb1a0402fced808dfcfd7b --- tensorflow/core/kernels/edit_distance_op.cc | 47 +++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tensorflow/core/kernels/edit_distance_op.cc b/tensorflow/core/kernels/edit_distance_op.cc index 4aecdc9e414d3..386a1af08409f 100644 --- a/tensorflow/core/kernels/edit_distance_op.cc +++ b/tensorflow/core/kernels/edit_distance_op.cc @@ -64,6 +64,12 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices, return errors::InvalidArgument( "truth_shape should be a vector, but got shape: ", truth_shape.shape().DebugString()); + if (hypothesis_values.NumElements() != hypothesis_indices.dim_size(0)) + return errors::InvalidArgument( + "Expected hypothesis_values.NumElements == " + "#rows(hypothesis_indices), their shapes are: ", + hypothesis_values.shape().DebugString(), " and ", + hypothesis_indices.shape().DebugString()); if (hypothesis_shape.NumElements() != hypothesis_indices.dim_size(1)) return errors::InvalidArgument( "Expected hypothesis_shape.NumElements == " @@ -75,6 +81,12 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices, "Input SparseTensors must have rank at least 2, but truth_shape " "rank is: ", truth_shape.NumElements()); + if (truth_values.NumElements() != truth_indices.dim_size(0)) + return errors::InvalidArgument( + "Expected truth_values.NumElements == " + "#rows(truth_indices), their shapes are: ", + truth_values.shape().DebugString(), " and ", + truth_indices.shape().DebugString()); if (truth_shape.NumElements() != truth_indices.dim_size(1)) return errors::InvalidArgument( "Expected truth_shape.NumElements == " @@ -153,6 +165,11 @@ class EditDistanceOp : public OpKernel { output_shape.AddDim(std::max(hypothesis_st_shape.dim_size(d), truth_st_shape.dim_size(d))); } + const auto output_elements = output_shape.num_elements(); + OP_REQUIRES( + ctx, output_elements > 0, + errors::InvalidArgument("Got output shape ", output_shape.DebugString(), + " which has 0 elements")); Tensor* output = nullptr; OP_REQUIRES_OK(ctx, ctx->allocate_output("output", output_shape, &output)); @@ -185,6 +202,12 @@ class EditDistanceOp : public OpKernel { if (g_truth == g_hypothesis) { auto loc = std::inner_product(g_truth.begin(), g_truth.end(), output_strides.begin(), int64{0}); + OP_REQUIRES( + ctx, loc < output_elements, + errors::Internal("Got an inner product ", loc, + " which would require in writing to outside of " + "the buffer for the output tensor (max elements ", + output_elements, ")")); output_t(loc) = gtl::LevenshteinDistance(truth_seq, hypothesis_seq, cmp); if (normalize_) output_t(loc) /= truth_seq.size(); @@ -194,6 +217,12 @@ class EditDistanceOp : public OpKernel { } else if (g_truth > g_hypothesis) { // zero-length truth auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(), output_strides.begin(), int64{0}); + OP_REQUIRES( + ctx, loc < output_elements, + errors::Internal("Got an inner product ", loc, + " which would require in writing to outside of " + "the buffer for the output tensor (max elements ", + output_elements, ")")); output_t(loc) = hypothesis_seq.size(); if (normalize_ && output_t(loc) != 0.0f) { output_t(loc) = std::numeric_limits::infinity(); @@ -202,6 +231,12 @@ class EditDistanceOp : public OpKernel { } else { // zero-length hypothesis auto loc = std::inner_product(g_truth.begin(), g_truth.end(), output_strides.begin(), int64{0}); + OP_REQUIRES( + ctx, loc < output_elements, + errors::Internal("Got an inner product ", loc, + " which would require in writing to outside of " + "the buffer for the output tensor (max elements ", + output_elements, ")")); output_t(loc) = (normalize_) ? 1.0 : truth_seq.size(); ++truth_iter; } @@ -212,6 +247,12 @@ class EditDistanceOp : public OpKernel { auto hypothesis_seq = hypothesis_j.values(); auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(), output_strides.begin(), int64{0}); + OP_REQUIRES( + ctx, loc < output_elements, + errors::Internal("Got an inner product ", loc, + " which would require in writing to outside of the " + "buffer for the output tensor (max elements ", + output_elements, ")")); output_t(loc) = hypothesis_seq.size(); if (normalize_ && output_t(loc) != 0.0f) { output_t(loc) = std::numeric_limits::infinity(); @@ -224,6 +265,12 @@ class EditDistanceOp : public OpKernel { auto truth_seq = truth_i.values(); auto loc = std::inner_product(g_truth.begin(), g_truth.end(), output_strides.begin(), int64{0}); + OP_REQUIRES( + ctx, loc < output_elements, + errors::Internal("Got an inner product ", loc, + " which would require in writing to outside of the " + "buffer for the output tensor (max elements ", + output_elements, ")")); output_t(loc) = (normalize_) ? 1.0 : truth_seq.size(); ++truth_iter; }