Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap OOB read in dequantize op.
Also fixes SEGV in same op

PiperOrigin-RevId: 372437896
Change-Id: I135e94d360c2a1ce374c10f7e0fed1af603dbc02
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 6, 2021
1 parent 2600359 commit 5899741
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/dequantize_op.cc
Expand Up @@ -98,6 +98,18 @@ class DequantizeOp : public OpKernel {
if (axis_ > -1) {
num_slices = input.dim_size(axis_);
}
OP_REQUIRES(ctx, input_min_tensor.NumElements() == num_slices,
errors::InvalidArgument(
"input_min_tensor must have as many elements as input on "
"the dequantization axis (",
axis_, "), got ", input_min_tensor.NumElements(),
", expected ", num_slices));
OP_REQUIRES(ctx, input_max_tensor.NumElements() == num_slices,
errors::InvalidArgument(
"input_max_tensor must have as many elements as input on "
"the dequantization axis (",
axis_, "), got ", input_max_tensor.NumElements(),
", expected ", num_slices));

Tensor* output = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));
Expand Down

0 comments on commit 5899741

Please sign in to comment.