Validate the rank and number of elements of the `num_bits` tensor for QuantizeAndDequantizeV3.upstream-r2.10-f3f9cb38ecf

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

2 files changed, 162 insertions, 105 deletions

diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
index da9257fb9c9..ae02b57861a 100644
--- a/tensorflow/core/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -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
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 21059a72e7c..a928f9ea4b4 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -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()