Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix Integer overflow error in Dequantize op shape function, by adding…
… a bound check on axis.

PiperOrigin-RevId: 412121389
Change-Id: I3088dbad9e90f9998d406b618c16694388a9dfb4
  • Loading branch information
ishark authored and tensorflower-gardener committed Nov 24, 2021
1 parent f3a46ea commit b64638e
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 0 deletions.
14 changes: 14 additions & 0 deletions tensorflow/core/ops/array_ops.cc
Expand Up @@ -24,6 +24,7 @@ limitations under the License.
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/mirror_pad_mode.h"
#include "tensorflow/core/util/padding.h"
#include "tensorflow/core/util/strided_slice_op.h"
Expand Down Expand Up @@ -3028,13 +3029,26 @@ REGISTER_OP("Dequantize")
return errors::InvalidArgument("axis should be at least -1, got ",
axis);
}
auto input_dims = c->Rank(c->input(0));
if (axis > input_dims) {
return errors::InvalidArgument(
"Axis must be less than input dimension(", input_dims, "), got ",
axis);
}
const int minmax_rank = (axis == -1) ? 0 : 1;
TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
ShapeHandle minmax;
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
if (axis != -1) {
ShapeHandle input;
if (axis >= kint32max) {
// Check int32 max bound for a corner case to prevent integer flow
// when input actually has kint32max rank and above bound check is not
// triggered.
return errors::InvalidArgument(
"Axis cannot be >= kint32max value, got ", axis);
}
TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
DimensionHandle depth;
TF_RETURN_IF_ERROR(
Expand Down
15 changes: 15 additions & 0 deletions tensorflow/python/kernel_tests/array_ops/array_ops_test.py
Expand Up @@ -1704,6 +1704,21 @@ def f(a):
output_grad = gradient_checker_v2.compute_gradient(f, [input_tensor])
self.assertAllClose(output_grad[0], np.zeros([1, 4, 4]))

def testOutOfBoundAxis(self):
input_tensor = constant_op.constant([1., 1.])
input_min = [0]
input_max = [1]
q_input, _, _ = array_ops.quantize(input_tensor, 0, 1, dtypes.qint32)
error = (errors.InvalidArgumentError, ValueError)
with self.assertRaisesRegex(error,
r".*Axis must be less than input dimension.*"):
self.evaluate(
gen_array_ops.dequantize(
input=q_input,
min_range=input_min,
max_range=input_max,
axis=2**31 - 1))


@test_util.run_all_in_graph_and_eager_modes
class SortedSearchTest(test_util.TensorFlowTestCase):
Expand Down

0 comments on commit b64638e

Please sign in to comment.