CVE-2022-35971: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. If `FakeQuantWithMinMaxVars` is given `min` or `max` tensors of a nonzero rank, it results in a `CHECK` fail that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 785d67a78a1d533759fcd2f5e8d6ef778de849e0. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
CVSS v3 base score: 7.5
Published: 2022-09-16T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/fake_quant_ops.cc, tensorflow/core/kernels/quantized_bias_add_op.cc, tensorflow/core/kernels/quantized_bias_add_op_test.cc, tensorflow/core/kernels/quantized_instance_norm.cc, tensorflow/core/kernels/requantize.cc, tensorflow/core/kernels/requantize_op_test.cc, tensorflow/python/kernel_tests/quantization_ops/BUILD, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-18 22:44:32 -0700
Commit message: Fix quantize ops input validation issues.

The majority of these are just missing checks on min/max.

PiperOrigin-RevId: 461800665

Diff: diff --git a/tensorflow/core/kernels/fake_quant_ops.cc b/tensorflow/core/kernels/fake_quant_ops.cc
index aa59213c67d81..b8c31701a61c0 100644
--- a/tensorflow/core/kernels/fake_quant_ops.cc
+++ b/tensorflow/core/kernels/fake_quant_ops.cc
@@ -22,10 +22,11 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/fake_quant_ops_functor.h"
 // Above is the related header but clang tidy doesn't recognize it.
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/monitoring/gauge.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/util/determinism.h"
 
@@ -203,10 +204,17 @@ class FakeQuantWithMinMaxVarsOp : public OpKernel {
     CHECK_EQ(3, context->num_inputs());
     const Tensor& input = context->input(0);
     const Tensor& min = context->input(1);
     const Tensor& max = context->input(2);
 
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(min.shape()),
+        InvalidArgument("`min` must be rank 0 but is rank ", min.dims()));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(max.shape()),
+        InvalidArgument("`max` must be rank 0 but is rank ", max.dims()));
+
     Tensor* output;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &output));
 
     FakeQuantWithMinMaxVarsFunctor<Device> functor;
@@ -340,14 +348,21 @@ class FakeQuantWithMinMaxVarsPerChannelOp : public OpKernel {
   void Compute(OpKernelContext* context) override {
     CHECK_EQ(3, context->num_inputs());
     const Tensor& input = context->input(0);
     const int depth = input.dim_size(input.dims() - 1);  // last dimension size.
     const Tensor& min = context->input(1);
+    const Tensor& max = context->input(2);
+
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVector(min.shape()),
+        InvalidArgument("`min` must be rank 1 but is rank ", min.dims()));
     OP_REQUIRES(context, min.dim_size(0) == depth,
                 InvalidArgument("min has incorrect size, expected ", depth,
                                 " was ", min.dim_size(0)));
-    const Tensor& max = context->input(2);
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVector(max.shape()),
+        InvalidArgument("`max` must be rank 1 but is rank ", max.dims()));
     OP_REQUIRES(context, max.dim_size(0) == depth,
                 InvalidArgument("max has incorrect size, expected ", depth,
                                 " was ", max.dim_size(0)));
 
     Tensor* output;
diff --git a/tensorflow/core/kernels/quantized_bias_add_op.cc b/tensorflow/core/kernels/quantized_bias_add_op.cc
index db0e21a498011..c064f9b1b21e2 100644
--- a/tensorflow/core/kernels/quantized_bias_add_op.cc
+++ b/tensorflow/core/kernels/quantized_bias_add_op.cc
@@ -18,10 +18,11 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/meta_support.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/kernels/quantization_utils.h"
 #include "tensorflow/core/lib/core/errors.h"
 
@@ -36,14 +37,34 @@ class QuantizedBiasAddOp : public OpKernel {
       : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
     const Tensor& bias = context->input(1);
-    const float input_min = context->input(2).flat<float>()(0);
-    const float input_max = context->input(3).flat<float>()(0);
-    const float bias_min = context->input(4).flat<float>()(0);
-    const float bias_max = context->input(5).flat<float>()(0);
+
+    const Tensor& min_input = context->input(2);
+    const Tensor& max_input = context->input(3);
+    const Tensor& min_bias = context->input(4);
+    const Tensor& max_bias = context->input(5);
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(min_input.shape()),
+        errors::InvalidArgument("`min_input` must be rank 0 but is rank ",
+                                min_input.dims()));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(max_input.shape()),
+        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
+                                max_input.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_bias.shape()),
+                errors::InvalidArgument(
+                    "`min_bias` must be rank 0 but is rank ", min_bias.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_bias.shape()),
+                errors::InvalidArgument(
+                    "`max_bias` must be rank 0 but is rank ", max_bias.dims()));
+
+    const float input_min = min_input.flat<float>()(0);
+    const float input_max = max_input.flat<float>()(0);
+    const float bias_min = min_bias.flat<float>()(0);
+    const float bias_max = max_bias.flat<float>()(0);
 
     OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input.shape()),
                 errors::InvalidArgument("Input tensor must be at least 2D: ",
                                         input.shape().DebugString()));
     OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
diff --git a/tensorflow/core/kernels/quantized_bias_add_op_test.cc b/tensorflow/core/kernels/quantized_bias_add_op_test.cc
index 7b99ceafe261b..edfae98efa953 100644
--- a/tensorflow/core/kernels/quantized_bias_add_op_test.cc
+++ b/tensorflow/core/kernels/quantized_bias_add_op_test.cc
@@ -72,14 +72,14 @@ TEST_F(QuantizedBiasAddTest, Small) {
 
   AddInputFromArray<quint8>(input_quantized.shape(),
                             input_quantized.flat<quint8>());
   AddInputFromArray<quint8>(bias_quantized.shape(),
                             bias_quantized.flat<quint8>());
-  AddInputFromArray<float>(TensorShape({1}), {input_min});
-  AddInputFromArray<float>(TensorShape({1}), {input_max});
-  AddInputFromArray<float>(TensorShape({1}), {bias_min});
-  AddInputFromArray<float>(TensorShape({1}), {bias_max});
+  AddInputFromArray<float>(TensorShape({}), {input_min});
+  AddInputFromArray<float>(TensorShape({}), {input_max});
+  AddInputFromArray<float>(TensorShape({}), {bias_min});
+  AddInputFromArray<float>(TensorShape({}), {bias_max});
   TF_ASSERT_OK(RunOpKernel());
   const Tensor& output_quantized = *GetOutput(0);
   const float output_min = GetOutput(1)->flat<float>()(0);
   const float output_max = GetOutput(2)->flat<float>()(0);
   Tensor output_float =
@@ -154,14 +154,14 @@ TEST_F(QuantizedBiasAddTest, RealData) {
 
   AddInputFromArray<quint8>(input_quantized.shape(),
                             input_quantized.flat<quint8>());
   AddInputFromArray<quint8>(bias_quantized.shape(),
                             bias_quantized.flat<quint8>());
-  AddInputFromArray<float>(TensorShape({1}), {input_min});
-  AddInputFromArray<float>(TensorShape({1}), {input_max});
-  AddInputFromArray<float>(TensorShape({1}), {bias_min});
-  AddInputFromArray<float>(TensorShape({1}), {bias_max});
+  AddInputFromArray<float>(TensorShape({}), {input_min});
+  AddInputFromArray<float>(TensorShape({}), {input_max});
+  AddInputFromArray<float>(TensorShape({}), {bias_min});
+  AddInputFromArray<float>(TensorShape({}), {bias_max});
   TF_ASSERT_OK(RunOpKernel());
   const Tensor& output_quantized = *GetOutput(0);
   const float output_min = GetOutput(1)->flat<float>()(0);
   const float output_max = GetOutput(2)->flat<float>()(0);
   Tensor output_float =
diff --git a/tensorflow/core/kernels/quantized_instance_norm.cc b/tensorflow/core/kernels/quantized_instance_norm.cc
index 4fd967fe3fc46..405c6aa95bb3e 100644
--- a/tensorflow/core/kernels/quantized_instance_norm.cc
+++ b/tensorflow/core/kernels/quantized_instance_norm.cc
@@ -23,11 +23,11 @@ limitations under the License.
 #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/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
-
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/quantization_utils.h"
 
 #ifdef USE_NEON
 namespace {
 
@@ -272,12 +272,20 @@ class QuantizedInstanceNorm : public OpKernel {
   }
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
 
-    float input_min = context->input(1).flat<float>()(0);
-    float input_max = context->input(2).flat<float>()(0);
+    const Tensor& x_min = context->input(1);
+    const Tensor& x_max = context->input(2);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_min.shape()),
+                errors::InvalidArgument("`x_min` must be rank 0 but is rank ",
+                                        x_min.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_max.shape()),
+                errors::InvalidArgument("`x_max` must be rank 0 but is rank ",
+                                        x_max.dims()));
+    float input_min = x_min.scalar<float>()();
+    float input_max = x_max.scalar<float>()();
     float input_scale = (input_max - input_min) / 255.0f;
 
     OP_REQUIRES(context, input_min < input_max,
                 errors::InvalidArgument(
                     "input_min must be less than input_max : ", input_min,
diff --git a/tensorflow/core/kernels/requantize.cc b/tensorflow/core/kernels/requantize.cc
index 3259e5ddd096a..bc5de17163926 100644
--- a/tensorflow/core/kernels/requantize.cc
+++ b/tensorflow/core/kernels/requantize.cc
@@ -16,13 +16,15 @@ limitations under the License.
 // See docs in ../ops/array_ops.cc.
 
 #define EIGEN_USE_THREADS
 
 #include <math.h>
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/type_traits.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/meta_support.h"
 #include "tensorflow/core/kernels/quantization_utils.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -36,14 +38,38 @@ class RequantizeOp : public OpKernel {
  public:
   explicit RequantizeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& input = ctx->input(0);
-    const float input_min_float = ctx->input(1).flat<float>()(0);
-    const float input_max_float = ctx->input(2).flat<float>()(0);
-    const float requested_output_min_float = ctx->input(3).flat<float>()(0);
-    const float requested_output_max_float = ctx->input(4).flat<float>()(0);
+
+    const Tensor& input_min = ctx->input(1);
+    const Tensor& input_max = ctx->input(2);
+    const Tensor& requested_output_min = ctx->input(3);
+    const Tensor& requested_output_max = ctx->input(4);
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsScalar(input_min.shape()),
+        errors::InvalidArgument("`input_min` must be rank 0 but is rank ",
+                                input_min.dims()));
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsScalar(input_max.shape()),
+        errors::InvalidArgument("`input_max` must be rank 0 but is rank ",
+                                input_max.dims()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_min.shape()),
+                errors::InvalidArgument(
+                    "`requested_output_min` must be rank 0 but is rank ",
+                    requested_output_min.dims()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_max.shape()),
+                errors::InvalidArgument(
+                    "`requested_output_max` must be rank 0 but is rank ",
+                    requested_output_max.dims()));
+
+    const float input_min_float = input_min.flat<float>()(0);
+    const float input_max_float = input_max.flat<float>()(0);
+    const float requested_output_min_float =
+        requested_output_min.flat<float>()(0);
+    const float requested_output_max_float =
+        requested_output_max.flat<float>()(0);
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));
     Tensor* output_min = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &output_min));
diff --git a/tensorflow/core/kernels/requantize_op_test.cc b/tensorflow/core/kernels/requantize_op_test.cc
index 133f92b029d8d..5663520fdbfda 100644
--- a/tensorflow/core/kernels/requantize_op_test.cc
+++ b/tensorflow/core/kernels/requantize_op_test.cc
@@ -51,14 +51,14 @@ TEST_F(RequantizeTest, HandCraftedRequantize) {
   const int value_count = 3;
 
   // Requantize to -1 to 1.
   AddInputFromArray<qint32>(TensorShape({value_count}),
                             {-(1 << 23), 0, (1 << 23)});
-  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {-1.0f});
-  AddInputFromArray<float>(TensorShape({1}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {-256.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
+  AddInputFromArray<float>(TensorShape({}), {-1.0f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
   TF_ASSERT_OK(RunOpKernel());
   Tensor expected(allocator(), DT_QUINT8, TensorShape({value_count}));
   test::FillValues<quint8>(&expected, {0, 128, 255});
   test::ExpectTensorEqual<quint8>(expected, *GetOutput(0));
   test::ExpectTensorEqual<float>(test::AsScalar<float>(-1.0f), *GetOutput(1));
@@ -69,28 +69,28 @@ TEST_F(RequantizeTest, InvalidOutputMin) {
   ConfigureRequantize();
   const int value_count = 3;
 
   AddInputFromArray<qint32>(TensorShape({value_count}),
                             {-(1 << 23), 0, (1 << 23)});
-  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {0.01f});
-  AddInputFromArray<float>(TensorShape({1}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {-256.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
+  AddInputFromArray<float>(TensorShape({}), {0.01f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
   EXPECT_EQ("requested_output_min must be <= 0, but got 0.01",
             RunOpKernel().error_message());
 }
 
 TEST_F(RequantizeTest, InvalidOutputMax) {
   ConfigureRequantize();
   const int value_count = 3;
 
   AddInputFromArray<qint32>(TensorShape({value_count}),
                             {-(1 << 23), 0, (1 << 23)});
-  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {-10.0f});
-  AddInputFromArray<float>(TensorShape({1}), {-11.0f});
+  AddInputFromArray<float>(TensorShape({}), {-256.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
+  AddInputFromArray<float>(TensorShape({}), {-10.0f});
+  AddInputFromArray<float>(TensorShape({}), {-11.0f});
   EXPECT_EQ(
       "requested_output_max must be >= requested_output_min, but got -11 and "
       "-10",
       RunOpKernel().error_message());
 }
diff --git a/tensorflow/python/kernel_tests/quantization_ops/BUILD b/tensorflow/python/kernel_tests/quantization_ops/BUILD
new file mode 100644
index 0000000000000..ff0be9898c601
--- /dev/null
+++ b/tensorflow/python/kernel_tests/quantization_ops/BUILD
@@ -0,0 +1,24 @@
+# Tests of TensorFlow quantization ops written using the Python API.
+
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
+
+package(
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+tf_py_test(
+    name = "quantization_ops_test",
+    size = "small",
+    srcs = ["quantization_ops_test.py"],
+    deps = [
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:math_ops",
+        "//third_party/py/numpy",
+    ],
+)
diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
new file mode 100644
index 0000000000000..af7cecf2ffb02
--- /dev/null
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -0,0 +1,210 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tf.quantize ops."""
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import googletest
+
+
+class FakeQuantWithMinMaxVarsOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars(
+              inputs=inputs, min=0.0, max=[[1.0], [2.0], [4.0]]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars(
+              inputs=inputs, min=[[1.0], [2.0], [4.0]], max=1.0))
+
+
+class FakeQuantWithMinMaxVarsPerChannelOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 1"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel(
+              inputs=inputs, min=[[0.0]], max=[1.0]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Dimensions must be equal|incorrect size"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel(
+              inputs=inputs, min=[0.0, 0.1], max=[1.0]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 1"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel(
+              inputs=inputs, min=[1.0], max=[[1.0]]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Dimensions must be equal|incorrect size"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel(
+              inputs=inputs, min=[0.0], max=[1.0, 1.1]))
+
+
+class QuantizedBiasedAddTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.qint8)
+    bias = constant_op.constant(np.int8(0), shape=[3], dtype=dtypes.qint8)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_bias_add(
+              input=inputs,
+              bias=bias,
+              min_input=[],
+              max_input=1.0,
+              min_bias=0.0,
+              max_bias=1.0,
+              out_type=dtypes.qint32))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_bias_add(
+              input=inputs,
+              bias=bias,
+              min_input=0.0,
+              max_input=[],
+              min_bias=0.0,
+              max_bias=1.0,
+              out_type=dtypes.qint32))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_bias_add(
+              input=inputs,
+              bias=bias,
+              min_input=0.0,
+              max_input=1.0,
+              min_bias=[],
+              max_bias=1.0,
+              out_type=dtypes.qint32))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_bias_add(
+              input=inputs,
+              bias=bias,
+              min_input=0.0,
+              max_input=1.0,
+              min_bias=0.0,
+              max_bias=[],
+              out_type=dtypes.qint32))
+
+
+class QuantizedInstanceNormOpTest(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)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          array_ops.quantized_instance_norm(
+              x=inputs, x_min=0.0, x_max=[[1.0], [2.0], [4.0]]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          array_ops.quantized_instance_norm(
+              x=inputs, x_min=[[1.0], [2.0], [4.0]], x_max=1.0))
+
+
+class RequantizeOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.int32(0), shape=[3, 3, 3, 3], dtype=dtypes.qint32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.requantize(
+              input=inputs,
+              input_min=[],
+              input_max=1.0,
+              requested_output_min=0.0,
+              requested_output_max=1.0,
+              out_type=dtypes.qint8))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.requantize(
+              input=inputs,
+              input_min=0.0,
+              input_max=[],
+              requested_output_min=0.0,
+              requested_output_max=1.0,
+              out_type=dtypes.qint8))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.requantize(
+              input=inputs,
+              input_min=0.0,
+              input_max=1.0,
+              requested_output_min=[],
+              requested_output_max=1.0,
+              out_type=dtypes.qint8))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.requantize(
+              input=inputs,
+              input_min=0.0,
+              input_max=1.0,
+              requested_output_min=0.0,
+              requested_output_max=[],
+              out_type=dtypes.qint8))
+
+
+if __name__ == "__main__":
+  googletest.main()

Vulnerable code:     const Tensor& max = context->input(2);
    const float input_min = context->input(2).flat<float>()(0);
    const float input_max = context->input(3).flat<float>()(0);
    const float bias_min = context->input(4).flat<float>()(0);
    const float bias_max = context->input(5).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
  AddInputFromArray<float>(TensorShape({1}), {bias_min});
  AddInputFromArray<float>(TensorShape({1}), {bias_max});
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
  AddInputFromArray<float>(TensorShape({1}), {bias_min});
  AddInputFromArray<float>(TensorShape({1}), {bias_max});

    float input_min = context->input(1).flat<float>()(0);
    float input_max = context->input(2).flat<float>()(0);
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
    const float input_min_float = ctx->input(1).flat<float>()(0);
    const float input_max_float = ctx->input(2).flat<float>()(0);
    const float requested_output_min_float = ctx->input(3).flat<float>()(0);
    const float requested_output_max_float = ctx->input(4).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {-1.0f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {0.01f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {-10.0f});
  AddInputFromArray<float>(TensorShape({1}), {-11.0f});
Fixed code: #include "tensorflow/core/framework/tensor_shape.h"
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(min.shape()),
        InvalidArgument("`min` must be rank 0 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max.shape()),
        InvalidArgument("`max` must be rank 0 but is rank ", max.dims()));

    const Tensor& max = context->input(2);

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(min.shape()),
        InvalidArgument("`min` must be rank 1 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(max.shape()),
        InvalidArgument("`max` must be rank 1 but is rank ", max.dims()));
#include "tensorflow/core/framework/tensor_shape.h"

    const Tensor& min_input = context->input(2);
    const Tensor& max_input = context->input(3);
    const Tensor& min_bias = context->input(4);
    const Tensor& max_bias = context->input(5);
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(min_input.shape()),
        errors::InvalidArgument("`min_input` must be rank 0 but is rank ",
                                min_input.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max_input.shape()),
        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
                                max_input.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_bias.shape()),
                errors::InvalidArgument(
                    "`min_bias` must be rank 0 but is rank ", min_bias.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_bias.shape()),
                errors::InvalidArgument(
                    "`max_bias` must be rank 0 but is rank ", max_bias.dims()));

    const float input_min = min_input.flat<float>()(0);
    const float input_max = max_input.flat<float>()(0);
    const float bias_min = min_bias.flat<float>()(0);
    const float bias_max = max_bias.flat<float>()(0);
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {bias_min});
  AddInputFromArray<float>(TensorShape({}), {bias_max});
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {bias_min});
  AddInputFromArray<float>(TensorShape({}), {bias_max});
#include "tensorflow/core/framework/tensor_shape.h"
    const Tensor& x_min = context->input(1);
    const Tensor& x_max = context->input(2);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_min.shape()),
                errors::InvalidArgument("`x_min` must be rank 0 but is rank ",
                                        x_min.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_max.shape()),
                errors::InvalidArgument("`x_max` must be rank 0 but is rank ",
                                        x_max.dims()));
    float input_min = x_min.scalar<float>()();
    float input_max = x_max.scalar<float>()();

#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"

    const Tensor& input_min = ctx->input(1);
    const Tensor& input_max = ctx->input(2);
    const Tensor& requested_output_min = ctx->input(3);
    const Tensor& requested_output_max = ctx->input(4);
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(input_min.shape()),
        errors::InvalidArgument("`input_min` must be rank 0 but is rank ",
                                input_min.dims()));
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(input_max.shape()),
        errors::InvalidArgument("`input_max` must be rank 0 but is rank ",
                                input_max.dims()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_min.shape()),
                errors::InvalidArgument(
                    "`requested_output_min` must be rank 0 but is rank ",
                    requested_output_min.dims()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_max.shape()),
                errors::InvalidArgument(
                    "`requested_output_max` must be rank 0 but is rank ",
                    requested_output_max.dims()));

    const float input_min_float = input_min.flat<float>()(0);
    const float input_max_float = input_max.flat<float>()(0);
    const float requested_output_min_float =
        requested_output_min.flat<float>()(0);
    const float requested_output_max_float =
        requested_output_max.flat<float>()(0);
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-1.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {0.01f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-10.0f});
  AddInputFromArray<float>(TensorShape({}), {-11.0f});
# Tests of TensorFlow quantization ops written using the Python API.

# buildifier: disable=same-origin-load
load("//tensorflow:tensorflow.bzl", "tf_py_test")

package(
    default_visibility = ["//tensorflow:internal"],
    licenses = ["notice"],
)

tf_py_test(
    name = "quantization_ops_test",
    size = "small",
    srcs = ["quantization_ops_test.py"],
    deps = [
        "//tensorflow/python:array_ops",
        "//tensorflow/python:client",
        "//tensorflow/python:client_testlib",
        "//tensorflow/python:framework",
        "//tensorflow/python:framework_for_generated_wrappers",
        "//tensorflow/python:math_ops",
        "//third_party/py/numpy",
    ],
)
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for tf.quantize ops."""
import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import googletest


class FakeQuantWithMinMaxVarsOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars(
              inputs=inputs, min=0.0, max=[[1.0], [2.0], [4.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars(
              inputs=inputs, min=[[1.0], [2.0], [4.0]], max=1.0))


class FakeQuantWithMinMaxVarsPerChannelOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[[0.0]], max=[1.0]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Dimensions must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[0.0, 0.1], max=[1.0]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[1.0], max=[[1.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Dimensions must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[0.0], max=[1.0, 1.1]))


class QuantizedBiasedAddTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.qint8)
    bias = constant_op.constant(np.int8(0), shape=[3], dtype=dtypes.qint8)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=[],
              max_input=1.0,
              min_bias=0.0,
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=[],
              min_bias=0.0,
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=1.0,
              min_bias=[],
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=1.0,
              min_bias=0.0,
              max_bias=[],
              out_type=dtypes.qint32))


class QuantizedInstanceNormOpTest(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)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.quantized_instance_norm(
              x=inputs, x_min=0.0, x_max=[[1.0], [2.0], [4.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.quantized_instance_norm(
              x=inputs, x_min=[[1.0], [2.0], [4.0]], x_max=1.0))


class RequantizeOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.int32(0), shape=[3, 3, 3, 3], dtype=dtypes.qint32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=[],
              input_max=1.0,
              requested_output_min=0.0,
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=[],
              requested_output_min=0.0,
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=1.0,
              requested_output_min=[],
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=1.0,
              requested_output_min=0.0,
              requested_output_max=[],
              out_type=dtypes.qint8))


if __name__ == "__main__":
  googletest.main()