Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix QuantizedAvgPool invalid rank issue.
Passing in an invalid rank in eager mode led to a runtime crash.
Updated the runtime error to match the one generated by the
shape function for better consistency.

PiperOrigin-RevId: 462419362
  • Loading branch information
cantonios authored and tensorflower-gardener committed Jul 21, 2022
1 parent 177036a commit 7cdf9d4
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 22 deletions.
48 changes: 30 additions & 18 deletions tensorflow/core/kernels/quantized_pooling_ops.cc
Expand Up @@ -15,18 +15,18 @@ limitations under the License.

// See docs in ../ops/nn_ops.cc.

#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/platform/errors.h"
#define EIGEN_USE_THREADS

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/kernels/pooling_ops_common.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/padding.h"
#include "tensorflow/core/util/tensor_format.h"
Expand Down Expand Up @@ -67,8 +67,20 @@ class QuantizedAvgPoolingOp : public OpKernel {
return;
}

const float min_input = context->input(1).flat<float>()(0);
const float max_input = context->input(2).flat<float>()(0);
const Tensor& min_input_tensor = context->input(1);
const Tensor& max_input_tensor = context->input(2);
OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
errors::InvalidArgument(
"min_input shape must be rank 0 but is rank ",
min_input_tensor.dims(),
", received shape: ", min_input_tensor.shape()));
OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
errors::InvalidArgument(
"max_input shape must be rank 0 but is rank ",
max_input_tensor.dims(),
", received shape: ", max_input_tensor.shape()));
const float min_input = context->input(1).scalar<float>()();
const float max_input = context->input(2).scalar<float>()();

OP_REQUIRES(context, params.depth_window == 1,
errors::Unimplemented("Non-spatial pooling is not "
Expand Down Expand Up @@ -119,20 +131,20 @@ class QuantizedMaxPoolingOp : public MaxPoolingOp<Device, T> {
: MaxPoolingOp<Device, T>(context) {}

void Compute(OpKernelContext* context) override {
auto min_input_tensor = context->input(1);
auto max_input_tensor = context->input(2);
OP_REQUIRES(
context, min_input_tensor.NumElements() == 1,
errors::InvalidArgument(
"min_input must be a scalar float value, got tensor with shape ",
min_input_tensor.shape()));
OP_REQUIRES(
context, max_input_tensor.NumElements() == 1,
errors::InvalidArgument(
"max_input must be a scalar float value, got tensor with shape ",
max_input_tensor.shape()));
const float min_input = context->input(1).flat<float>()(0);
const float max_input = context->input(2).flat<float>()(0);
const Tensor& min_input_tensor = context->input(1);
const Tensor& max_input_tensor = context->input(2);
OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
errors::InvalidArgument(
"min_input shape must be rank 0 but is rank ",
min_input_tensor.dims(),
", received shape: ", min_input_tensor.shape()));
OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
errors::InvalidArgument(
"max_input shape must be rank 0 but is rank ",
max_input_tensor.dims(),
", received shape: ", max_input_tensor.shape()));
const float min_input = context->input(1).scalar<float>()();
const float max_input = context->input(2).scalar<float>()();
MaxPoolingOp<Device, T>::Compute(context);
Tensor* output_min = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min));
Expand Down
8 changes: 4 additions & 4 deletions tensorflow/core/kernels/quantized_pooling_ops_test.cc
Expand Up @@ -69,8 +69,8 @@ TEST_F(QuantizedPoolingTest, SmallAveragePooling) {

AddInputFromArray<quint8>(input_quantized.shape(),
input_quantized.flat<quint8>());
AddInputFromArray<float>(TensorShape({1}), {input_min});
AddInputFromArray<float>(TensorShape({1}), {input_max});
AddInputFromArray<float>(TensorShape({}), {input_min});
AddInputFromArray<float>(TensorShape({}), {input_max});
TF_ASSERT_OK(RunOpKernel());
const Tensor& output_quantized = *GetOutput(0);
const float output_min = GetOutput(1)->flat<float>()(0);
Expand Down Expand Up @@ -114,8 +114,8 @@ TEST_F(QuantizedPoolingTest, SmallMaxPooling) {

AddInputFromArray<quint8>(input_quantized.shape(),
input_quantized.flat<quint8>());
AddInputFromArray<float>(TensorShape({1}), {input_min});
AddInputFromArray<float>(TensorShape({1}), {input_max});
AddInputFromArray<float>(TensorShape({}), {input_min});
AddInputFromArray<float>(TensorShape({}), {input_max});
TF_ASSERT_OK(RunOpKernel());
const Tensor& output_quantized = *GetOutput(0);
const float output_min = GetOutput(1)->flat<float>()(0);
Expand Down
Expand Up @@ -154,6 +154,72 @@ def test_invalid_inputs(self):
x=inputs, x_min=[[1.0], [2.0], [4.0]], x_max=1.0))


class QuantizedAvgPoolingOpTest(test_util.TensorFlowTestCase):

@test_util.run_in_graph_and_eager_modes
def test_invalid_inputs(self):
inputs = constant_op.constant(
np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
ksize = [1, 1, 1, 1]
strides = [1, 1, 1, 1]
padding = "SAME"

with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
"must be.* rank 0"):
self.evaluate(
nn_ops.quantized_avg_pool(
input=inputs,
min_input=[],
max_input=1.0,
ksize=ksize,
strides=strides,
padding=padding))

with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
"must be.* rank 0"):
self.evaluate(
nn_ops.quantized_avg_pool(
input=inputs,
min_input=0.0,
max_input=[],
ksize=ksize,
strides=strides,
padding=padding))


class QuantizedMaxPoolingOpTest(test_util.TensorFlowTestCase):

@test_util.run_in_graph_and_eager_modes
def test_invalid_inputs(self):
inputs = constant_op.constant(
np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
ksize = [1, 1, 1, 1]
strides = [1, 1, 1, 1]
padding = "SAME"

with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
"must be.* rank 0"):
self.evaluate(
nn_ops.quantized_max_pool(
input=inputs,
min_input=[],
max_input=1.0,
ksize=ksize,
strides=strides,
padding=padding))

with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
"must be.* rank 0"):
self.evaluate(
nn_ops.quantized_max_pool(
input=inputs,
min_input=0.0,
max_input=[],
ksize=ksize,
strides=strides,
padding=padding))


class RequantizeOpTest(test_util.TensorFlowTestCase):

@test_util.run_in_graph_and_eager_modes
Expand Down

1 comment on commit 7cdf9d4

@elfringham
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is causing a unit test failure on --config=mkl_aarch64
#56861

Please sign in to comment.