43 lines
2.1 KiB
Diff
43 lines
2.1 KiB
Diff
From efea03b38fb8d3b81762237dc85e579cc5fc6e87 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
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<float>()(0);
|
|
- const float max_x = context->input(3).flat<float>()(0);
|
|
- const float min_y = context->input(4).flat<float>()(0);
|
|
- const float max_y = context->input(5).flat<float>()(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<float>()(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<float>()(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<float>()(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<float>()(0);
|
|
|
|
BCast bcast(BCast::FromShape(x.shape()), BCast::FromShape(y.shape()));
|
|
if (!bcast.IsValid()) {
|