From efea03b38fb8d3b81762237dc85e579cc5fc6e87 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Wed, 21 Apr 2021 16:15:46 -0700 Subject: [PATCH] Validate inputs to `QuantizedMul` PiperOrigin-RevId: 369756982 Change-Id: I00d960cc3b9316fd7a86bd37a44e341c96e17624 --- tensorflow/core/kernels/quantized_mul_op.cc | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/kernels/quantized_mul_op.cc b/tensorflow/core/kernels/quantized_mul_op.cc index fb56f68bf14db..22cff8939449a 100644 --- a/tensorflow/core/kernels/quantized_mul_op.cc +++ b/tensorflow/core/kernels/quantized_mul_op.cc @@ -284,10 +284,22 @@ class QuantizedMulOp : public OpKernel { void Compute(OpKernelContext* context) override { const Tensor& x = context->input(0); const Tensor& y = context->input(1); - const float min_x = context->input(2).flat()(0); - const float max_x = context->input(3).flat()(0); - const float min_y = context->input(4).flat()(0); - const float max_y = context->input(5).flat()(0); + auto& min_x_tensor = context->input(2); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_x_tensor.shape()), + errors::InvalidArgument("min_x must be a scalar")); + const float min_x = min_x_tensor.flat()(0); + auto& max_x_tensor = context->input(3); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_x_tensor.shape()), + errors::InvalidArgument("max_x must be a scalar")); + const float max_x = max_x_tensor.flat()(0); + auto& min_y_tensor = context->input(4); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_y_tensor.shape()), + errors::InvalidArgument("min_y must be a scalar")); + const float min_y = min_y_tensor.flat()(0); + auto& max_y_tensor = context->input(5); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_y_tensor.shape()), + errors::InvalidArgument("max_y must be a scalar")); + const float max_y = max_y_tensor.flat()(0); BCast bcast(BCast::FromShape(x.shape()), BCast::FromShape(y.shape())); if (!bcast.IsValid()) {