CVE-2022-36026 (published 2022-09-16T22:15Z)

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.

CVSS v3 base score: 7.5, HIGH (CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: CWE-617
cwe_name: Reachable Assertion
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.
extended_description: 
            While assertion is good for catching logic errors and reducing the chances of reaching more serious vulnerability conditions, it can still lead to a denial of service.
            For example, if a server handles multiple simultaneous connections, and an assert() occurs in one single connection that causes all other connections to be dropped, this is a reachable assertion that leads to a denial of service.
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Dan Suh
author_date: 2022-07-27 22:22:30-07:00
msg: 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
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: quantize_and_dequantize_op.cc
change_type: ModificationType.MODIFY
diff: @@ -21,19 +21,23 @@ limitations under the License.
 #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
@@ -49,8 +53,8 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
     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;
@@ -58,10 +62,10 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
     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") {
@@ -72,12 +76,10 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
 
   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;
@@ -91,21 +93,21 @@ class QuantizeAndDequantizeV2Op : public OpKernel {
         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});
@@ -158,38 +160,34 @@ class QuantizeAndDequantizeV4GradientOp : public OpKernel {
     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)));
     }
 
@@ -203,12 +201,12 @@ class QuantizeAndDequantizeV4GradientOp : public OpKernel {
                    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>(),
@@ -252,21 +250,25 @@ 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;
@@ -274,24 +276,24 @@ class QuantizeAndDequantizeV3Op : public OpKernel {
       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});
@@ -331,15 +333,14 @@ class QuantizeAndDequantizeOp : public OpKernel {
     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_));
     }
   }
@@ -371,53 +372,53 @@ class QuantizeAndDequantizeOp : public OpKernel {
   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,
@@ -425,16 +426,16 @@ struct QuantizeAndDequantizePerChannelGradientFunctor<CPUDevice, T> {
                   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
 
@@ -442,22 +443,22 @@ template struct functor::QuantizeAndDequantizePerChannelGradientFunctor<
   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
@@ -470,29 +471,29 @@ TF_CALL_double(REGISTER_CPU_KERNEL);
                               .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

programming_language: C++
filename: quantization_ops_test.py
change_type: ModificationType.MODIFY
diff: @@ -15,9 +15,11 @@
 """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
@@ -407,5 +409,59 @@ class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):
               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()

programming_language: Python