Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix multiple issues in EditDistance
PiperOrigin-RevId: 372033948
Change-Id: Ieb957c29894af05bdfeb1a0402fced808dfcfd7b
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 5, 2021
1 parent 79865b5 commit f4c364a
Showing 1 changed file with 47 additions and 0 deletions.
47 changes: 47 additions & 0 deletions tensorflow/core/kernels/edit_distance_op.cc
Expand Up @@ -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 == "
Expand All @@ -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 == "
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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<T>(truth_seq, hypothesis_seq, cmp);
if (normalize_) output_t(loc) /= truth_seq.size();
Expand All @@ -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<float>::infinity();
Expand All @@ -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;
}
Expand All @@ -212,6 +247,12 @@ class EditDistanceOp : public OpKernel {
auto hypothesis_seq = hypothesis_j.values<T>();
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<float>::infinity();
Expand All @@ -224,6 +265,12 @@ class EditDistanceOp : public OpKernel {
auto truth_seq = truth_i.values<T>();
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;
}
Expand Down

0 comments on commit f4c364a

Please sign in to comment.