Visitar URL original
Merge pull request #58570 from tensorflow/r2.9-7b174a0f2e4 · tensorflow/tensorflow@5dbe90a · GitHub
Skip to content

Commit 5dbe90a

Browse files
mihaimaruseacmihaimaruseac
authored
Merge pull request #58570 from tensorflow/r2.9-7b174a0f2e4
r2.9 cherry-pick: 7b174a0 "Fix asan issue with QuantizeAndDequantizeV2/V3/V4/V4Grad shape inference functions."
2 parents fb6df70 + eb92930 commit 5dbe90a

2 files changed

Lines changed: 82 additions & 0 deletions

File tree

‎tensorflow/core/ops/array_ops.cc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2879,6 +2879,10 @@ REGISTER_OP("QuantizeAndDequantizeV2")
28792879
axis);
28802880
} else if (axis != -1) {
28812881
ShapeHandle input;
2882+
if (axis >= kint32max) {
2883+
return errors::InvalidArgument(
2884+
"Axis cannot be >= kint32max value, got ", axis);
2885+
}
28822886
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
28832887
DimensionHandle depth;
28842888
TF_RETURN_IF_ERROR(
@@ -2914,6 +2918,10 @@ REGISTER_OP("QuantizeAndDequantizeV4")
29142918
axis);
29152919
} else if (axis != -1) {
29162920
ShapeHandle input;
2921+
if (axis >= kint32max) {
2922+
return errors::InvalidArgument(
2923+
"Axis cannot be >= kint32max value, got ", axis);
2924+
}
29172925
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29182926
DimensionHandle depth;
29192927
TF_RETURN_IF_ERROR(
@@ -2945,6 +2953,10 @@ REGISTER_OP("QuantizeAndDequantizeV4Grad")
29452953
axis);
29462954
} else if (axis != -1) {
29472955
ShapeHandle input;
2956+
if (axis >= kint32max) {
2957+
return errors::InvalidArgument(
2958+
"Axis cannot be >= kint32max value, got ", axis);
2959+
}
29482960
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29492961
DimensionHandle depth;
29502962
TF_RETURN_IF_ERROR(
@@ -2981,6 +2993,10 @@ REGISTER_OP("QuantizeAndDequantizeV3")
29812993
axis);
29822994
} else if (axis != -1) {
29832995
ShapeHandle input;
2996+
if (axis >= kint32max) {
2997+
return errors::InvalidArgument(
2998+
"Axis cannot be >= kint32max value, got ", axis);
2999+
}
29843000
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
29853001
DimensionHandle depth;
29863002
TF_RETURN_IF_ERROR(

‎tensorflow/python/kernel_tests/array_ops/array_ops_test.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1791,6 +1791,72 @@ def testOutOfBoundAxis(self):
17911791
max_range=input_max,
17921792
axis=2**31 - 1))
17931793

1794+
@test_util.run_v2_only
1795+
def testInvalidAxis(self):
1796+
1797+
@def_function.function
1798+
def test_quantize_and_dequantize_v2():
1799+
gen_array_ops.quantize_and_dequantize_v2(
1800+
input=[2.5],
1801+
input_min=[1.0],
1802+
input_max=[10.0],
1803+
signed_input=True,
1804+
num_bits=1,
1805+
range_given=True,
1806+
round_mode="HALF_TO_EVEN",
1807+
narrow_range=True,
1808+
axis=0x7fffffff)
1809+
1810+
@def_function.function
1811+
def test_quantize_and_dequantize_v3():
1812+
gen_array_ops.quantize_and_dequantize_v3(
1813+
input=[2.5],
1814+
input_min=[1.0],
1815+
input_max=[10.0],
1816+
num_bits=1,
1817+
signed_input=True,
1818+
range_given=True,
1819+
narrow_range=True,
1820+
axis=0x7fffffff)
1821+
1822+
@def_function.function
1823+
def test_quantize_and_dequantize_v4():
1824+
gen_array_ops.quantize_and_dequantize_v4(
1825+
input=[2.5],
1826+
input_min=[1.0],
1827+
input_max=[10.0],
1828+
signed_input=True,
1829+
num_bits=1,
1830+
range_given=True,
1831+
round_mode="HALF_TO_EVEN",
1832+
narrow_range=True,
1833+
axis=0x7fffffff)
1834+
1835+
@def_function.function
1836+
def test_quantize_and_dequantize_v4_grad():
1837+
gen_array_ops.quantize_and_dequantize_v4_grad(
1838+
gradients=[2.5],
1839+
input=[2.5],
1840+
input_min=[1.0],
1841+
input_max=[10.0],
1842+
axis=0x7fffffff)
1843+
1844+
with self.assertRaisesRegex(
1845+
ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
1846+
test_quantize_and_dequantize_v2()
1847+
1848+
with self.assertRaisesRegex(
1849+
ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
1850+
test_quantize_and_dequantize_v3()
1851+
1852+
with self.assertRaisesRegex(
1853+
ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
1854+
test_quantize_and_dequantize_v4()
1855+
1856+
with self.assertRaisesRegex(
1857+
ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
1858+
test_quantize_and_dequantize_v4_grad()
1859+
17941860

17951861
@test_util.run_all_in_graph_and_eager_modes
17961862
class SortedSearchTest(test_util.TensorFlowTestCase):

0 commit comments

Comments
 (0)