Skip to content

Commit 7cf73a2

Browse files
pkanwar23tensorflower-gardener
authored andcommitted
Address QuantizeAndDequantizeV* heap oob. Added additional checks for the 'axis' attribute.
PiperOrigin-RevId: 402446942 Change-Id: Id2f6b82e4e740d0550329be02621c46466b5a5b9
1 parent 3bbc2ee commit 7cf73a2

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

Diff for: tensorflow/core/ops/array_ops.cc

+16-4
Original file line numberDiff line numberDiff line change
@@ -2863,7 +2863,10 @@ REGISTER_OP("QuantizeAndDequantizeV2")
28632863
ShapeHandle minmax;
28642864
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
28652865
TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
2866-
if (axis != -1) {
2866+
if (axis < -1) {
2867+
return errors::InvalidArgument("axis should be at least -1, got ",
2868+
axis);
2869+
} else if (axis != -1) {
28672870
ShapeHandle input;
28682871
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
28692872
DimensionHandle depth;
@@ -2895,7 +2898,10 @@ REGISTER_OP("QuantizeAndDequantizeV4")
28952898
ShapeHandle minmax;
28962899
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
28972900
TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
2898-
if (axis != -1) {
2901+
if (axis < -1) {
2902+
return errors::InvalidArgument("axis should be at least -1, got ",
2903+
axis);
2904+
} else if (axis != -1) {
28992905
ShapeHandle input;
29002906
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29012907
DimensionHandle depth;
@@ -2923,7 +2929,10 @@ REGISTER_OP("QuantizeAndDequantizeV4Grad")
29232929
ShapeHandle minmax;
29242930
TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
29252931
TF_RETURN_IF_ERROR(c->Merge(c->input(3), minmax, &minmax));
2926-
if (axis != -1) {
2932+
if (axis < -1) {
2933+
return errors::InvalidArgument("axis should be at least -1, got ",
2934+
axis);
2935+
} else if (axis != -1) {
29272936
ShapeHandle input;
29282937
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29292938
DimensionHandle depth;
@@ -2956,7 +2965,10 @@ REGISTER_OP("QuantizeAndDequantizeV3")
29562965
ShapeHandle minmax;
29572966
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
29582967
TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
2959-
if (axis != -1) {
2968+
if (axis < -1) {
2969+
return errors::InvalidArgument("axis should be at least -1, got ",
2970+
axis);
2971+
} else if (axis != -1) {
29602972
ShapeHandle input;
29612973
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29622974
DimensionHandle depth;

Diff for: tensorflow/core/ops/array_ops_test.cc

+2
Original file line numberDiff line numberDiff line change
@@ -1374,6 +1374,8 @@ TEST(ArrayOpsTest, QuantizeAndDequantizeV2_ShapeFn) {
13741374
INFER_ERROR("Shapes must be equal rank, but are 1 and 0", op,
13751375
"[1,2,?,4,5];[];[1]");
13761376
INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[1];[1]");
1377+
(*op.node_def.mutable_attr())["axis"].set_i(-2);
1378+
INFER_ERROR("axis should be at least -1, got -2", op, "?;?;?");
13771379
}
13781380

13791381
TEST(ArrayOpsTest, SpaceToBatch_ShapeFn) {

0 commit comments

Comments
 (0)