From c5b0d5f8ac19888e46ca14b0e27562e7fbbee9a9 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Fri, 30 Apr 2021 10:39:05 -0700
Subject: [PATCH] Fix the CHECK failure in tf.raw_ops.QuantizeAndDequantizeV2.

---
 tensorflow/core/kernels/quantize_and_dequantize_op.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
index ae905e22..408196a2 100644
--- a/tensorflow/core/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -72,6 +72,9 @@ 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,
-- 
2.27.0