From 8a793b5d7f59e37ac7f3cd0954a750a2fe76bad4 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Thu, 29 Jul 2021 18:23:45 -0700 Subject: [PATCH] Prevent division by 0 in common shape functions. PiperOrigin-RevId: 387712197 Change-Id: Id25c7460e35b68aeeeac23b9a88e455b443ee149 --- tensorflow/core/framework/common_shape_fns.cc | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc index b9efddf4..e25d5581 100644 --- a/tensorflow/core/framework/common_shape_fns.cc +++ b/tensorflow/core/framework/common_shape_fns.cc @@ -659,6 +659,8 @@ Status Conv2DShapeImpl(shape_inference::InferenceContext* c, if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) { int64 input_depth_value = c->Value(input_depth_dim), filter_input_depth_value = c->Value(filter_input_depth_dim); + if (filter_input_depth_value == 0) + return errors::InvalidArgument("Depth of filter must not be 0"); if (input_depth_value % filter_input_depth_value != 0) return errors::InvalidArgument( "Depth of input (", input_depth_value, @@ -668,6 +670,8 @@ Status Conv2DShapeImpl(shape_inference::InferenceContext* c, int64 num_groups = input_depth_value / filter_input_depth_value; if (c->ValueKnown(output_depth_dim)) { int64 output_depth_value = c->Value(output_depth_dim); + if (num_groups == 0) + return errors::InvalidArgument("Number of groups must not be 0"); if (output_depth_value % num_groups != 0) return errors::InvalidArgument( "Depth of output (", output_depth_value, @@ -798,6 +802,8 @@ Status Conv3DShape(shape_inference::InferenceContext* c) { if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) { int64 input_depth_value = c->Value(input_depth_dim), filter_input_depth_value = c->Value(filter_input_depth_dim); + if (filter_input_depth_value == 0) + return errors::InvalidArgument("Depth of filter must not be 0"); if (input_depth_value % filter_input_depth_value != 0) return errors::InvalidArgument( "Depth of input (", input_depth_value, @@ -807,6 +813,8 @@ Status Conv3DShape(shape_inference::InferenceContext* c) { int64 num_groups = input_depth_value / filter_input_depth_value; if (c->ValueKnown(output_depth_dim)) { int64 output_depth_value = c->Value(output_depth_dim); + if (num_groups == 0) + return errors::InvalidArgument("Number of groups must not be 0"); if (output_depth_value % num_groups != 0) return errors::InvalidArgument( "Depth of output (", output_depth_value, @@ -2364,6 +2372,9 @@ Status SparseReduceShapeFn(InferenceContext* c) { int64 ndims = shape_vec.size(); absl::flat_hash_set axes; + if (ndims == 0) + return errors::InvalidArgument( + "Number of dims in shape tensor must not be 0"); for (int i = 0; i < axes_vec.size(); i++) { axes.insert((axes_vec(i) + ndims) % ndims); } -- 2.27.0