CVE-2022-36026: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. If `QuantizeAndDequantizeV3` is given a nonscalar `num_bits` input tensor, 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 f3f9cb38ecfe5a8a703f2c4a8fead434ef291713. 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-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/quantize_and_dequantize_op.cc, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-27 22:25:32 -0700
Commit message: Validate the rank and number of elements of the `num_bits` tensor for QuantizeAndDequantizeV3.

QuantizeAndDequantizeV3Op, which accepts `num_bits` as a tensor, has a precondition that it should be rank <= 1 and the number of elements should be 1.
This change adds a validation for the Compute() method for this condition.

PiperOrigin-RevId: 463755293

Diff: diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
index da9257fb9c9af..ae02b57861ac0 100644
--- a/tensorflow/core/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -19,23 +19,27 @@ limitations under the License.
 #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
     (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
 #define EIGEN_USE_GPU
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
-#include "tensorflow/core/kernels/quantize_and_dequantize_op.h"
-
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.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/quantize_and_dequantize_op.h"
 #include "tensorflow/core/lib/core/errors.h"
 
 namespace tensorflow {
+namespace {
 
-typedef Eigen::ThreadPoolDevice CPUDevice;
-typedef Eigen::GpuDevice GPUDevice;
+using CpuDevice = ::Eigen::ThreadPoolDevice;
+using GpuDevice = ::Eigen::GpuDevice;
+using ::tensorflow::errors::InvalidArgument;
+
+}  // namespace
 
 // Simulate quantization precision loss in a float tensor by:
 // 1. Quantize the tensor to fixed point numbers, which should match the target
 //    quantization method when it is used in inference.
 // 2. Dequantize it back to floating point numbers for the following ops, most
@@ -47,39 +51,37 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
       : OpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("signed_input", &signed_input_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &axis_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
     OP_REQUIRES(ctx, num_bits_ > 0 && num_bits_ < (signed_input_ ? 62 : 63),
-                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
-                                        " with signed_input_ ", signed_input_));
+                InvalidArgument("num_bits is out of range: ", num_bits_,
+                                " with signed_input_ ", signed_input_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("range_given", &range_given_));
 
     string round_mode_string;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("round_mode", &round_mode_string));
     OP_REQUIRES(
         ctx,
         (round_mode_string == "HALF_UP" || round_mode_string == "HALF_TO_EVEN"),
-        errors::InvalidArgument("Round mode string must be "
-                                "'HALF_UP' or "
-                                "'HALF_TO_EVEN', is '" +
-                                round_mode_string + "'"));
+        InvalidArgument("Round mode string must be "
+                        "'HALF_UP' or "
+                        "'HALF_TO_EVEN', is '" +
+                        round_mode_string + "'"));
     if (round_mode_string == "HALF_UP") {
       round_mode_ = ROUND_HALF_UP;
     } else if (round_mode_string == "HALF_TO_EVEN") {
       round_mode_ = ROUND_HALF_TO_EVEN;
     }
     OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range_));
   }
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& input = ctx->input(0);
-    OP_REQUIRES(
-        ctx, axis_ >= -1,
-        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
-    OP_REQUIRES(
-        ctx, (axis_ == -1 || axis_ < input.shape().dims()),
-        errors::InvalidArgument("Shape must be at least rank ", axis_ + 1,
+    OP_REQUIRES(ctx, axis_ >= -1,
+                InvalidArgument("Axis must be at least -1. Found ", axis_));
+    OP_REQUIRES(ctx, (axis_ == -1 || axis_ < input.shape().dims()),
+                InvalidArgument("Shape must be at least rank ", axis_ + 1,
                                 " but is rank ", input.shape().dims()));
     const int depth = (axis_ == -1) ? 1 : input.dim_size(axis_);
     Tensor input_min_tensor;
     Tensor input_max_tensor;
     Tensor* output = nullptr;
@@ -89,25 +91,25 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
       input_max_tensor = ctx->input(2);
       if (axis_ == -1) {
         auto min_val = input_min_tensor.scalar<T>()();
         auto max_val = input_max_tensor.scalar<T>()();
         OP_REQUIRES(ctx, min_val <= max_val,
-                    errors::InvalidArgument("Invalid range: input_min ",
-                                            min_val, " > input_max ", max_val));
+                    InvalidArgument("Invalid range: input_min ", min_val,
+                                    " > input_max ", max_val));
       } else {
-        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
-                    errors::InvalidArgument(
-                        "input_min_tensor has incorrect size, was ",
-                        input_min_tensor.dim_size(0), " expected ", depth,
-                        " to match dim ", axis_, " of the input ",
-                        input_min_tensor.shape()));
-        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
-                    errors::InvalidArgument(
-                        "input_max_tensor has incorrect size, was ",
-                        input_max_tensor.dim_size(0), " expected ", depth,
-                        " to match dim ", axis_, " of the input ",
-                        input_max_tensor.shape()));
+        OP_REQUIRES(
+            ctx, input_min_tensor.dim_size(0) == depth,
+            InvalidArgument("input_min_tensor has incorrect size, was ",
+                            input_min_tensor.dim_size(0), " expected ", depth,
+                            " to match dim ", axis_, " of the input ",
+                            input_min_tensor.shape()));
+        OP_REQUIRES(
+            ctx, input_max_tensor.dim_size(0) == depth,
+            InvalidArgument("input_max_tensor has incorrect size, was ",
+                            input_max_tensor.dim_size(0), " expected ", depth,
+                            " to match dim ", axis_, " of the input ",
+                            input_max_tensor.shape()));
       }
     } else {
       auto range_shape = (axis_ == -1) ? TensorShape({}) : TensorShape({depth});
       OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                              range_shape, &input_min_tensor));
@@ -156,42 +158,38 @@ class QuantizeAndDequantizeV4GradientOp : public OpKernel {
     const Tensor& gradient = ctx->input(0);
     const Tensor& input = ctx->input(1);
     Tensor* input_backprop = nullptr;
     OP_REQUIRES_OK(ctx,
                    ctx->allocate_output(0, input.shape(), &input_backprop));
-    OP_REQUIRES(
-        ctx, axis_ >= -1,
-        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
+    OP_REQUIRES(ctx, axis_ >= -1,
+                InvalidArgument("Axis must be at least -1. Found ", axis_));
     OP_REQUIRES(ctx, (axis_ == -1 || axis_ < input.shape().dims()),
-                errors::InvalidArgument(
+                InvalidArgument(
                     "Axis should be -1 or 0 or a positive value less than ",
                     input.shape().dims(), "but given axis value was ", axis_));
 
-    OP_REQUIRES(
-        ctx, input.IsSameSize(gradient),
-        errors::InvalidArgument("gradient and input must be the same size"));
+    OP_REQUIRES(ctx, input.IsSameSize(gradient),
+                InvalidArgument("gradient and input must be the same size"));
     const int depth = (axis_ == -1) ? 1 : input.dim_size(axis_);
     const Tensor& input_min_tensor = ctx->input(2);
     OP_REQUIRES(ctx,
                 input_min_tensor.dims() == 0 || input_min_tensor.dims() == 1,
-                errors::InvalidArgument(
+                InvalidArgument(
                     "Input min tensor must have dimension 0 or 1. Received ",
                     input_min_tensor.dims(), "."));
     const Tensor& input_max_tensor = ctx->input(3);
     OP_REQUIRES(ctx,
                 input_max_tensor.dims() == 0 || input_max_tensor.dims() == 1,
-                errors::InvalidArgument(
+                InvalidArgument(
                     "Input max tensor must have dimension 0 or 1. Received ",
                     input_max_tensor.dims(), "."));
     if (axis_ != -1) {
-      OP_REQUIRES(
-          ctx, input_min_tensor.dim_size(0) == depth,
-          errors::InvalidArgument("min has incorrect size, expected ", depth,
+      OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
+                  InvalidArgument("min has incorrect size, expected ", depth,
                                   " was ", input_min_tensor.dim_size(0)));
-      OP_REQUIRES(
-          ctx, input_max_tensor.dim_size(0) == depth,
-          errors::InvalidArgument("max has incorrect size, expected ", depth,
+      OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
+                  InvalidArgument("max has incorrect size, expected ", depth,
                                   " was ", input_max_tensor.dim_size(0)));
     }
 
     TensorShape min_max_shape(input_min_tensor.shape());
     Tensor* input_min_backprop;
@@ -201,16 +199,16 @@ class QuantizeAndDequantizeV4GradientOp : public OpKernel {
     Tensor* input_max_backprop;
     OP_REQUIRES_OK(ctx,
                    ctx->allocate_output(2, min_max_shape, &input_max_backprop));
 
     if (axis_ == -1) {
-      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
-                  errors::InvalidArgument(
-                      "input_min must be a scalar if axis is unspecified"));
-      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
-                  errors::InvalidArgument(
-                      "input_max must be a scalar if axis is unspecified"));
+      OP_REQUIRES(
+          ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
+          InvalidArgument("input_min must be a scalar if axis is unspecified"));
+      OP_REQUIRES(
+          ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
+          InvalidArgument("input_max must be a scalar if axis is unspecified"));
       functor::QuantizeAndDequantizeOneScaleGradientFunctor<Device, T> f;
       f(ctx->eigen_device<Device>(), gradient.template flat<T>(),
         input.template flat<T>(), input_min_tensor.scalar<T>(),
         input_max_tensor.scalar<T>(), input_backprop->template flat<T>(),
         input_min_backprop->template scalar<T>(),
@@ -250,50 +248,54 @@ class QuantizeAndDequantizeV3Op : public OpKernel {
   }
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& input = ctx->input(0);
     OP_REQUIRES(ctx, axis_ < input.dims(),
-                errors::InvalidArgument(
+                InvalidArgument(
                     "Axis requested is larger than input dimensions. Axis: ",
                     axis_, " Input Dimensions: ", input.dims()));
     const int depth = (axis_ == -1) ? 1 : input.dim_size(axis_);
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));
 
-    Tensor num_bits_tensor;
-    num_bits_tensor = ctx->input(3);
-    int num_bits_val = num_bits_tensor.scalar<int32>()();
+    // Get num_bits and validate.
+    const Tensor num_bits_tensor = ctx->input(3);
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(num_bits_tensor.shape()),
+                InvalidArgument("Invalid shape. The `num_bits` tensor should "
+                                "be a scalar. Got dimensions: ",
+                                num_bits_tensor.dims()));
 
-    OP_REQUIRES(
-        ctx, num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
-        errors::InvalidArgument("num_bits is out of range: ", num_bits_val,
-                                " with signed_input_ ", signed_input_));
+    const int num_bits_val = num_bits_tensor.scalar<int32>()();
+    OP_REQUIRES(ctx,
+                num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
+                InvalidArgument("num_bits is out of range: ", num_bits_val,
+                                " with `signed_input_` ", signed_input_));
 
     Tensor input_min_tensor;
     Tensor input_max_tensor;
     if (range_given_) {
       input_min_tensor = ctx->input(1);
       input_max_tensor = ctx->input(2);
       if (axis_ == -1) {
-        auto min_val = input_min_tensor.scalar<T>()();
-        auto max_val = input_max_tensor.scalar<T>()();
+        const auto min_val = input_min_tensor.scalar<T>()();
+        const auto max_val = input_max_tensor.scalar<T>()();
         OP_REQUIRES(ctx, min_val <= max_val,
-                    errors::InvalidArgument("Invalid range: input_min ",
-                                            min_val, " > input_max ", max_val));
+                    InvalidArgument("Invalid range: input_min ", min_val,
+                                    " > input_max ", max_val));
       } else {
-        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
-                    errors::InvalidArgument(
-                        "input_min_tensor has incorrect size, was ",
-                        input_min_tensor.dim_size(0), " expected ", depth,
-                        " to match dim ", axis_, " of the input ",
-                        input_min_tensor.shape()));
-        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
-                    errors::InvalidArgument(
-                        "input_max_tensor has incorrect size, was ",
-                        input_max_tensor.dim_size(0), " expected ", depth,
-                        " to match dim ", axis_, " of the input ",
-                        input_max_tensor.shape()));
+        OP_REQUIRES(
+            ctx, input_min_tensor.dim_size(0) == depth,
+            InvalidArgument("input_min_tensor has incorrect size, was ",
+                            input_min_tensor.dim_size(0), " expected ", depth,
+                            " to match dim ", axis_, " of the input ",
+                            input_min_tensor.shape()));
+        OP_REQUIRES(
+            ctx, input_max_tensor.dim_size(0) == depth,
+            InvalidArgument("input_max_tensor has incorrect size, was ",
+                            input_max_tensor.dim_size(0), " expected ", depth,
+                            " to match dim ", axis_, " of the input ",
+                            input_max_tensor.shape()));
       }
     } else {
       auto range_shape = (axis_ == -1) ? TensorShape({}) : TensorShape({depth});
       OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                              range_shape, &input_min_tensor));
@@ -329,19 +331,18 @@ class QuantizeAndDequantizeOp : public OpKernel {
  public:
   explicit QuantizeAndDequantizeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("signed_input", &signed_input_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
     OP_REQUIRES(ctx, num_bits_ > 0 && num_bits_ < (signed_input_ ? 62 : 63),
-                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
-                                        " with signed_input_ ", signed_input_));
+                InvalidArgument("num_bits is out of range: ", num_bits_,
+                                " with signed_input_ ", signed_input_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("range_given", &range_given_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("input_min", &input_min_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("input_max", &input_max_));
     if (range_given_) {
-      OP_REQUIRES(
-          ctx, input_min_ <= input_max_,
-          errors::InvalidArgument("Invalid range: input_min ", input_min_,
+      OP_REQUIRES(ctx, input_min_ <= input_max_,
+                  InvalidArgument("Invalid range: input_min ", input_min_,
                                   " > input_max ", input_max_));
     }
   }
 
   void Compute(OpKernelContext* ctx) override {
@@ -369,97 +370,97 @@ class QuantizeAndDequantizeOp : public OpKernel {
   bool range_given_;
   float input_min_;
   float input_max_;
 };
 
-// Specializations for CPUDevice.
+// Specializations for CpuDevice.
 
 namespace functor {
 template <typename T>
-struct QuantizeAndDequantizeOneScaleFunctor<CPUDevice, T> {
-  void operator()(const CPUDevice& d, typename TTypes<T>::ConstVec input,
+struct QuantizeAndDequantizeOneScaleFunctor<CpuDevice, T> {
+  void operator()(const CpuDevice& d, typename TTypes<T>::ConstVec input,
                   const bool signed_input, const int num_bits,
                   const bool range_given, Tensor* input_min_tensor,
                   Tensor* input_max_tensor, QuantizerRoundMode round_mode,
                   bool narrow_range, typename TTypes<T>::Vec out) {
-    QuantizeAndDequantizeOneScaleImpl<CPUDevice, T>::Compute(
+    QuantizeAndDequantizeOneScaleImpl<CpuDevice, T>::Compute(
         d, input, signed_input, num_bits, range_given, input_min_tensor,
         input_max_tensor, round_mode, narrow_range, out);
   }
 };
 
 template <typename T>
-struct QuantizeAndDequantizePerChannelFunctor<CPUDevice, T> {
-  void operator()(const CPUDevice& d, typename TTypes<T, 3>::ConstTensor input,
+struct QuantizeAndDequantizePerChannelFunctor<CpuDevice, T> {
+  void operator()(const CpuDevice& d, typename TTypes<T, 3>::ConstTensor input,
                   bool signed_input, int num_bits, bool range_given,
                   Tensor* input_min_tensor, Tensor* input_max_tensor,
                   QuantizerRoundMode round_mode, bool narrow_range,
                   typename TTypes<T, 3>::Tensor out) {
-    QuantizeAndDequantizePerChannelImpl<CPUDevice, T>::Compute(
+    QuantizeAndDequantizePerChannelImpl<CpuDevice, T>::Compute(
         d, input, signed_input, num_bits, range_given, input_min_tensor,
         input_max_tensor, round_mode, narrow_range, out);
   }
 };
 
 template <typename T>
-struct QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice, T> {
-  void operator()(const CPUDevice& d, typename TTypes<T>::ConstFlat gradient,
+struct QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice, T> {
+  void operator()(const CpuDevice& d, typename TTypes<T>::ConstFlat gradient,
                   typename TTypes<T>::ConstFlat input,
                   typename TTypes<T>::ConstScalar input_min_tensor,
                   typename TTypes<T>::ConstScalar input_max_tensor,
                   typename TTypes<T>::Flat input_backprop,
                   typename TTypes<T>::Scalar input_min_backprop,
                   typename TTypes<T>::Scalar input_max_backprop) {
-    QuantizeAndDequantizeOneScaleGradientImpl<CPUDevice, T>::Compute(
+    QuantizeAndDequantizeOneScaleGradientImpl<CpuDevice, T>::Compute(
         d, gradient, input, input_min_tensor, input_max_tensor, input_backprop,
         input_min_backprop, input_max_backprop);
   }
 };
 
 template <typename T>
-struct QuantizeAndDequantizePerChannelGradientFunctor<CPUDevice, T> {
-  void operator()(const CPUDevice& d,
+struct QuantizeAndDequantizePerChannelGradientFunctor<CpuDevice, T> {
+  void operator()(const CpuDevice& d,
                   typename TTypes<T, 3>::ConstTensor gradient,
                   typename TTypes<T, 3>::ConstTensor input,
                   const Tensor* input_min_tensor,
                   const Tensor* input_max_tensor,
                   typename TTypes<T, 3>::Tensor input_backprop,
                   typename TTypes<T>::Flat input_min_backprop,
                   typename TTypes<T>::Flat input_max_backprop) {
-    QuantizeAndDequantizePerChannelGradientImpl<CPUDevice, T>::Compute(
+    QuantizeAndDequantizePerChannelGradientImpl<CpuDevice, T>::Compute(
         d, gradient, input, input_min_tensor, input_max_tensor, input_backprop,
         input_min_backprop, input_max_backprop);
   }
 };
 
-template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice,
+template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice,
                                                                       float>;
 template struct functor::QuantizeAndDequantizePerChannelGradientFunctor<
-    CPUDevice, double>;
+    CpuDevice, double>;
 
 }  // namespace functor
 
 #define REGISTER_CPU_KERNEL(T)                                                 \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV2")                      \
                               .Device(DEVICE_CPU)                              \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
+                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV3")                      \
                               .Device(DEVICE_CPU)                              \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV3Op<CPUDevice, T>);            \
+                          QuantizeAndDequantizeV3Op<CpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV4")                      \
                               .Device(DEVICE_CPU)                              \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
+                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV4Grad")                  \
                               .Device(DEVICE_CPU)                              \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV4GradientOp<CPUDevice, T>);    \
+                          QuantizeAndDequantizeV4GradientOp<CpuDevice, T>);    \
   REGISTER_KERNEL_BUILDER(                                                     \
       Name("QuantizeAndDequantize").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
-      QuantizeAndDequantizeOp<CPUDevice, T>);
+      QuantizeAndDequantizeOp<CpuDevice, T>);
 TF_CALL_float(REGISTER_CPU_KERNEL);
 TF_CALL_double(REGISTER_CPU_KERNEL);
 #undef REGISTER_CPU_KERNEL
 
 #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
@@ -468,33 +469,33 @@ TF_CALL_double(REGISTER_CPU_KERNEL);
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV2")                      \
                               .Device(DEVICE_GPU)                              \
                               .HostMemory("input_min")                         \
                               .HostMemory("input_max")                         \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
+                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV3")                      \
                               .Device(DEVICE_GPU)                              \
                               .HostMemory("input_min")                         \
                               .HostMemory("input_max")                         \
                               .HostMemory("num_bits")                          \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV3Op<GPUDevice, T>);            \
+                          QuantizeAndDequantizeV3Op<GpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV4")                      \
                               .Device(DEVICE_GPU)                              \
                               .HostMemory("input_min")                         \
                               .HostMemory("input_max")                         \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
+                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
   REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV4Grad")                  \
                               .Device(DEVICE_GPU)                              \
                               .HostMemory("input_min")                         \
                               .HostMemory("input_max")                         \
                               .TypeConstraint<T>("T"),                         \
-                          QuantizeAndDequantizeV4GradientOp<GPUDevice, T>);    \
+                          QuantizeAndDequantizeV4GradientOp<GpuDevice, T>);    \
   REGISTER_KERNEL_BUILDER(                                                     \
       Name("QuantizeAndDequantize").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
-      QuantizeAndDequantizeOp<GPUDevice, T>);
+      QuantizeAndDequantizeOp<GpuDevice, T>);
 TF_CALL_float(REGISTER_GPU_KERNEL);
 TF_CALL_double(REGISTER_GPU_KERNEL);
 #undef REGISTER_GPU_KERNEL
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 }  // namespace tensorflow
diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
index 21059a72e7cc9..a928f9ea4b4db 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -13,13 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for tf.quantize ops."""
 import numpy as np
 
+from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
 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
@@ -405,7 +407,61 @@ class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):
           math_ops.quantize_down_and_shrink_range(
               input=inputs, input_min=[], input_max=4.0,
               out_type=dtypes.quint8))
 
 
+class QuantizeAndDequantizeV3OpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_valid(self):
+    with ops.Graph().as_default(), context.eager_mode():
+      input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
+                                         shape=(6,),
+                                         dtype=dtypes.float32),
+      input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
+      input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
+      num_bits = constant_op.constant(8, shape=(), dtype=dtypes.int32)
+
+      quantized = array_ops.quantize_and_dequantize_v3(
+          input_value,
+          input_min,
+          input_max,
+          num_bits,
+          signed_input=True,
+          range_given=False)
+      self.assertSequenceAlmostEqual(
+          input_value[0].numpy(), quantized.numpy()[0], delta=0.05)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
+                                       shape=(6,),
+                                       dtype=dtypes.float32),
+    input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
+    input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
+    # Tensor with invalid shape and invalid number of elements.
+    num_bits = constant_op.constant([], shape=(0,), dtype=dtypes.int32)
+
+    # Test that running the op raises error. It raises different errors
+    # depending on whether the shape inference is run first or the op's
+    # Compute() is run first.
+    try:
+      array_ops.quantize_and_dequantize_v3(
+          input_value, input_min, input_max, num_bits, signed_input=True)
+    except Exception as ex:  # pylint: disable=broad-except
+      if isinstance(ex, errors.InvalidArgumentError):
+        self.assertRegex(str(ex), "The `num_bits` tensor should be a scalar.")
+      elif isinstance(ex, ValueError):
+        self.assertRegex(str(ex), "Shape must be rank 0")
+      else:
+        self.fail(
+            "Raised exception other than expected: %s. "
+            "Expected exceptions are errors.InvalidArgumentError or ValueError",
+            ex.__name__)
+    else:
+      self.fail(
+          "Did not raise an exception where it is expected to raise either "
+          "a ValueError or errors.InvalidArgumentError.")
+
+
 if __name__ == "__main__":
   googletest.main()

Vulnerable code: #include "tensorflow/core/kernels/quantize_and_dequantize_op.h"

typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
                                        " with signed_input_ ", signed_input_));
        errors::InvalidArgument("Round mode string must be "
                                "'HALF_UP' or "
                                "'HALF_TO_EVEN', is '" +
                                round_mode_string + "'"));
    OP_REQUIRES(
        ctx, axis_ >= -1,
        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
    OP_REQUIRES(
        ctx, (axis_ == -1 || axis_ < input.shape().dims()),
        errors::InvalidArgument("Shape must be at least rank ", axis_ + 1,
                    errors::InvalidArgument("Invalid range: input_min ",
                                            min_val, " > input_max ", max_val));
        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_min_tensor has incorrect size, was ",
                        input_min_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_min_tensor.shape()));
        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_max_tensor has incorrect size, was ",
                        input_max_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_max_tensor.shape()));
    OP_REQUIRES(
        ctx, axis_ >= -1,
        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
                errors::InvalidArgument(
    OP_REQUIRES(
        ctx, input.IsSameSize(gradient),
        errors::InvalidArgument("gradient and input must be the same size"));
                errors::InvalidArgument(
                errors::InvalidArgument(
      OP_REQUIRES(
          ctx, input_min_tensor.dim_size(0) == depth,
          errors::InvalidArgument("min has incorrect size, expected ", depth,
      OP_REQUIRES(
          ctx, input_max_tensor.dim_size(0) == depth,
          errors::InvalidArgument("max has incorrect size, expected ", depth,
      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
                  errors::InvalidArgument(
                      "input_min must be a scalar if axis is unspecified"));
      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
                  errors::InvalidArgument(
                      "input_max must be a scalar if axis is unspecified"));
                errors::InvalidArgument(
    Tensor num_bits_tensor;
    num_bits_tensor = ctx->input(3);
    int num_bits_val = num_bits_tensor.scalar<int32>()();
    OP_REQUIRES(
        ctx, num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
        errors::InvalidArgument("num_bits is out of range: ", num_bits_val,
                                " with signed_input_ ", signed_input_));
        auto min_val = input_min_tensor.scalar<T>()();
        auto max_val = input_max_tensor.scalar<T>()();
                    errors::InvalidArgument("Invalid range: input_min ",
                                            min_val, " > input_max ", max_val));
        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_min_tensor has incorrect size, was ",
                        input_min_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_min_tensor.shape()));
        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_max_tensor has incorrect size, was ",
                        input_max_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_max_tensor.shape()));
                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
                                        " with signed_input_ ", signed_input_));
      OP_REQUIRES(
          ctx, input_min_ <= input_max_,
          errors::InvalidArgument("Invalid range: input_min ", input_min_,
// Specializations for CPUDevice.
struct QuantizeAndDequantizeOneScaleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::ConstVec input,
    QuantizeAndDequantizeOneScaleImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T, 3>::ConstTensor input,
    QuantizeAndDequantizePerChannelImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::ConstFlat gradient,
    QuantizeAndDequantizeOneScaleGradientImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelGradientFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d,
    QuantizeAndDequantizePerChannelGradientImpl<CPUDevice, T>::Compute(
template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice,
    CPUDevice, double>;
                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV3Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<CPUDevice, T>);    \
      QuantizeAndDequantizeOp<CPUDevice, T>);
                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV3Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<GPUDevice, T>);    \
      QuantizeAndDequantizeOp<GPUDevice, T>);
Fixed code: #include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/quantize_and_dequantize_op.h"
namespace {
using CpuDevice = ::Eigen::ThreadPoolDevice;
using GpuDevice = ::Eigen::GpuDevice;
using ::tensorflow::errors::InvalidArgument;

}  // namespace
                InvalidArgument("num_bits is out of range: ", num_bits_,
                                " with signed_input_ ", signed_input_));
        InvalidArgument("Round mode string must be "
                        "'HALF_UP' or "
                        "'HALF_TO_EVEN', is '" +
                        round_mode_string + "'"));
    OP_REQUIRES(ctx, axis_ >= -1,
                InvalidArgument("Axis must be at least -1. Found ", axis_));
    OP_REQUIRES(ctx, (axis_ == -1 || axis_ < input.shape().dims()),
                InvalidArgument("Shape must be at least rank ", axis_ + 1,
                    InvalidArgument("Invalid range: input_min ", min_val,
                                    " > input_max ", max_val));
        OP_REQUIRES(
            ctx, input_min_tensor.dim_size(0) == depth,
            InvalidArgument("input_min_tensor has incorrect size, was ",
                            input_min_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_min_tensor.shape()));
        OP_REQUIRES(
            ctx, input_max_tensor.dim_size(0) == depth,
            InvalidArgument("input_max_tensor has incorrect size, was ",
                            input_max_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_max_tensor.shape()));
    OP_REQUIRES(ctx, axis_ >= -1,
                InvalidArgument("Axis must be at least -1. Found ", axis_));
                InvalidArgument(
    OP_REQUIRES(ctx, input.IsSameSize(gradient),
                InvalidArgument("gradient and input must be the same size"));
                InvalidArgument(
                InvalidArgument(
      OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                  InvalidArgument("min has incorrect size, expected ", depth,
      OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                  InvalidArgument("max has incorrect size, expected ", depth,
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
          InvalidArgument("input_min must be a scalar if axis is unspecified"));
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
          InvalidArgument("input_max must be a scalar if axis is unspecified"));
                InvalidArgument(
    // Get num_bits and validate.
    const Tensor num_bits_tensor = ctx->input(3);
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(num_bits_tensor.shape()),
                InvalidArgument("Invalid shape. The `num_bits` tensor should "
                                "be a scalar. Got dimensions: ",
                                num_bits_tensor.dims()));
    const int num_bits_val = num_bits_tensor.scalar<int32>()();
    OP_REQUIRES(ctx,
                num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
                InvalidArgument("num_bits is out of range: ", num_bits_val,
                                " with `signed_input_` ", signed_input_));
        const auto min_val = input_min_tensor.scalar<T>()();
        const auto max_val = input_max_tensor.scalar<T>()();
                    InvalidArgument("Invalid range: input_min ", min_val,
                                    " > input_max ", max_val));
        OP_REQUIRES(
            ctx, input_min_tensor.dim_size(0) == depth,
            InvalidArgument("input_min_tensor has incorrect size, was ",
                            input_min_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_min_tensor.shape()));
        OP_REQUIRES(
            ctx, input_max_tensor.dim_size(0) == depth,
            InvalidArgument("input_max_tensor has incorrect size, was ",
                            input_max_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_max_tensor.shape()));
                InvalidArgument("num_bits is out of range: ", num_bits_,
                                " with signed_input_ ", signed_input_));
      OP_REQUIRES(ctx, input_min_ <= input_max_,
                  InvalidArgument("Invalid range: input_min ", input_min_,
// Specializations for CpuDevice.
struct QuantizeAndDequantizeOneScaleFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T>::ConstVec input,
    QuantizeAndDequantizeOneScaleImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T, 3>::ConstTensor input,
    QuantizeAndDequantizePerChannelImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T>::ConstFlat gradient,
    QuantizeAndDequantizeOneScaleGradientImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelGradientFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d,
    QuantizeAndDequantizePerChannelGradientImpl<CpuDevice, T>::Compute(
template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice,
    CpuDevice, double>;
                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV3Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<CpuDevice, T>);    \
      QuantizeAndDequantizeOp<CpuDevice, T>);
                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV3Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<GpuDevice, T>);    \
      QuantizeAndDequantizeOp<GpuDevice, T>);
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
class QuantizeAndDequantizeV3OpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_valid(self):
    with ops.Graph().as_default(), context.eager_mode():
      input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
                                         shape=(6,),
                                         dtype=dtypes.float32),
      input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
      input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
      num_bits = constant_op.constant(8, shape=(), dtype=dtypes.int32)

      quantized = array_ops.quantize_and_dequantize_v3(
          input_value,
          input_min,
          input_max,
          num_bits,
          signed_input=True,
          range_given=False)
      self.assertSequenceAlmostEqual(
          input_value[0].numpy(), quantized.numpy()[0], delta=0.05)

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
                                       shape=(6,),
                                       dtype=dtypes.float32),
    input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
    input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
    # Tensor with invalid shape and invalid number of elements.
    num_bits = constant_op.constant([], shape=(0,), dtype=dtypes.int32)

    # Test that running the op raises error. It raises different errors
    # depending on whether the shape inference is run first or the op's
    # Compute() is run first.
    try:
      array_ops.quantize_and_dequantize_v3(
          input_value, input_min, input_max, num_bits, signed_input=True)
    except Exception as ex:  # pylint: disable=broad-except
      if isinstance(ex, errors.InvalidArgumentError):
        self.assertRegex(str(ex), "The `num_bits` tensor should be a scalar.")
      elif isinstance(ex, ValueError):
        self.assertRegex(str(ex), "Shape must be rank 0")
      else:
        self.fail(
            "Raised exception other than expected: %s. "
            "Expected exceptions are errors.InvalidArgumentError or ValueError",
            ex.__name__)
    else:
      self.fail(
          "Did not raise an exception where it is expected to raise either "
          "a ValueError or errors.InvalidArgumentError.")