Validate axis input in tf.raw_ops.QuantizeAndDequantizeV4Grad

PiperOrigin-RevId: 388291385 Change-Id: I3bab68dc61d935afa96c0da021a7b722c6dc8dc8
maciekpac · Aug 2, 2021 · 96f364a · 96f364a
1 parent 10fe168
commit 96f364a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/quantize_and_dequantize_op.cc b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -158,6 +158,13 @@ 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 || axis_ < input.shape().dims()),
+                errors::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),