From d6ed5bcfe1dcab9e85a4d39931bd18d99018e75b Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Fri, 23 Apr 2021 11:40:06 -0700 Subject: [PATCH] Add missing validation in `QuantizedBatchNormWithGlobalNormalization` PiperOrigin-RevId: 370123451 Change-Id: Id234d6dab1ec21230bb8e503dba30f899af87f33 --- .../core/kernels/quantized_batch_norm_op.cc | 77 ++++++++++++++++--- 1 file changed, 67 insertions(+), 10 deletions(-) diff --git a/tensorflow/core/kernels/quantized_batch_norm_op.cc b/tensorflow/core/kernels/quantized_batch_norm_op.cc index b03da7ad17fab..6dfe07f97a400 100644 --- a/tensorflow/core/kernels/quantized_batch_norm_op.cc +++ b/tensorflow/core/kernels/quantized_batch_norm_op.cc @@ -173,20 +173,50 @@ class QuantizedBatchNormOp : public OpKernel { void Compute(OpKernelContext* context) override { const Tensor& input = context->input(0); - const float input_min = context->input(1).flat()(0); - const float input_max = context->input(2).flat()(0); + const auto& input_min_tensor = context->input(1); + OP_REQUIRES(context, input_min_tensor.NumElements() == 1, + errors::InvalidArgument("input_min must have 1 element")); + const float input_min = input_min_tensor.flat()(0); + const auto& input_max_tensor = context->input(2); + OP_REQUIRES(context, input_max_tensor.NumElements() == 1, + errors::InvalidArgument("input_max must have 1 element")); + const float input_max = input_max_tensor.flat()(0); const Tensor& mean = context->input(3); - const float mean_min = context->input(4).flat()(0); - const float mean_max = context->input(5).flat()(0); + const auto& mean_min_tensor = context->input(4); + OP_REQUIRES(context, mean_min_tensor.NumElements() == 1, + errors::InvalidArgument("mean_min must have 1 element")); + const float mean_min = mean_min_tensor.flat()(0); + const auto& mean_max_tensor = context->input(5); + OP_REQUIRES(context, mean_max_tensor.NumElements() == 1, + errors::InvalidArgument("mean_max must have 1 element")); + const float mean_max = mean_max_tensor.flat()(0); const Tensor& var = context->input(6); - const float var_min = context->input(7).flat()(0); - const float var_max = context->input(8).flat()(0); + const auto& var_min_tensor = context->input(7); + OP_REQUIRES(context, var_min_tensor.NumElements() == 1, + errors::InvalidArgument("var_min must have 1 element")); + const float var_min = var_min_tensor.flat()(0); + const auto& var_max_tensor = context->input(8); + OP_REQUIRES(context, var_max_tensor.NumElements() == 1, + errors::InvalidArgument("var_max must have 1 element")); + const float var_max = var_max_tensor.flat()(0); const Tensor& beta = context->input(9); - const float beta_min = context->input(10).flat()(0); - const float beta_max = context->input(11).flat()(0); + const auto& beta_min_tensor = context->input(10); + OP_REQUIRES(context, beta_min_tensor.NumElements() == 1, + errors::InvalidArgument("beta_min must have 1 element")); + const float beta_min = beta_min_tensor.flat()(0); + const auto& beta_max_tensor = context->input(11); + OP_REQUIRES(context, beta_max_tensor.NumElements() == 1, + errors::InvalidArgument("beta_max must have 1 element")); + const float beta_max = beta_max_tensor.flat()(0); const Tensor& gamma = context->input(12); - const float gamma_min = context->input(13).flat()(0); - const float gamma_max = context->input(14).flat()(0); + const auto& gamma_min_tensor = context->input(13); + OP_REQUIRES(context, gamma_min_tensor.NumElements() == 1, + errors::InvalidArgument("gamma_min must have 1 element")); + const float gamma_min = gamma_min_tensor.flat()(0); + const auto& gamma_max_tensor = context->input(14); + OP_REQUIRES(context, gamma_max_tensor.NumElements() == 1, + errors::InvalidArgument("gamma_max must have 1 element")); + const float gamma_max = gamma_max_tensor.flat()(0); OP_REQUIRES(context, input.dims() == 4, errors::InvalidArgument("input must be 4-dimensional", @@ -203,6 +233,33 @@ class QuantizedBatchNormOp : public OpKernel { OP_REQUIRES(context, gamma.dims() == 1, errors::InvalidArgument("gamma must be 1-dimensional", gamma.shape().DebugString())); + OP_REQUIRES(context, mean.NumElements() > 1, + errors::InvalidArgument("Must have at least a mean value", + gamma.shape().DebugString())); + OP_REQUIRES(context, mean.NumElements() > 1, + errors::InvalidArgument("Must have at least a mean value")); + const auto last_dim = input.shape().dims() - 1; + OP_REQUIRES(context, + mean.shape().dim_size(0) == input.shape().dim_size(last_dim), + errors::InvalidArgument("Must provide as many means as the " + "last dimension of the input tensor: ", + mean.shape().DebugString(), " vs. ", + input.shape().DebugString())); + OP_REQUIRES( + context, mean.shape().dim_size(0) == var.shape().dim_size(0), + errors::InvalidArgument( + "Mean and variance tensors must have the same shape: ", + mean.shape().DebugString(), " vs. ", var.shape().DebugString())); + OP_REQUIRES( + context, mean.shape().dim_size(0) == beta.shape().dim_size(0), + errors::InvalidArgument( + "Mean and beta tensors must have the same shape: ", + mean.shape().DebugString(), " vs. ", beta.shape().DebugString())); + OP_REQUIRES( + context, mean.shape().dim_size(0) == gamma.shape().dim_size(0), + errors::InvalidArgument( + "Mean and gamma tensors must have the same shape: ", + mean.shape().DebugString(), " vs. ", gamma.shape().DebugString())); Tensor* output = nullptr; OP_REQUIRES_OK(context,