Merge pull request #58242 from tensorflow/r2.9-d66e1d56827

r2.9 cherry-pick: d66e1d5 "Fix tensor shape overflow in FusedResizeAndPadConv2D."
tensorflow · Oct 24, 2022 · 987d6f6 · 987d6f6
2 parents 340ad8e + 73a8cff
commit 987d6f6
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git a/tensorflow/core/kernels/conv_ops_fused_image_transform.cc b/tensorflow/core/kernels/conv_ops_fused_image_transform.cc
@@ -667,8 +667,11 @@ class FusedResizeConv2DUsingGemmOp : public OpKernel {
       st.height_scale = 1.0f;
       st.width_scale = 1.0f;
     }
-    TensorShape resized_shape(
-        {input.dim_size(0), st.out_height, st.out_width, input.dim_size(3)});
+    TensorShape resized_shape;
+    OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
+                                {input.dim_size(0), st.out_height, st.out_width,
+                                 input.dim_size(3)},
+                                &resized_shape));
     int paddings_index;
     int filter_index;
     if (DoResize) {

diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
@@ -574,7 +574,7 @@ REGISTER_OP("FusedResizeAndPadConv2D")
     .Attr("strides: list(int)")
     .Attr(GetPaddingAttrString())
     .SetShapeFn([](InferenceContext* c) {
-      return CommonFusedConvCalculations(c, true /* has_resize */);
+      return CommonFusedConvCalculations(c, /*has_resize=*/true);
     });
 
 REGISTER_OP("FusedPadConv2D")
@@ -587,7 +587,7 @@ REGISTER_OP("FusedPadConv2D")
     .Attr("strides: list(int)")
     .Attr(GetPaddingAttrString())
     .SetShapeFn([](InferenceContext* c) {
-      return CommonFusedConvCalculations(c, false /* has_resize */);
+      return CommonFusedConvCalculations(c, /*has_resize=*/false);
     });
 
 // --------------------------------------------------------------------------

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3431,6 +3431,33 @@ def testAddWithSameSrcAndAddTensorBuffer(self):
         np.rint(expected_output),
         self.evaluate(add).reshape(-1))
 
+  # Fused resize and pad conv.
+  @test_util.run_in_graph_and_eager_modes()
+  def testResizeAndPadLargeResize(self):
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                "Encountered overflow"):
+      mode = "REFLECT"
+      strides = [1, 1, 1, 1]
+      padding = "SAME"
+      resize_align_corners = False
+      tensor = constant_op.constant(
+          147, shape=[3, 3, 1, 4], dtype=dtypes.float32)
+      size = constant_op.constant([1879048192, 1879048192], dtype=dtypes.int32)
+      paddings = constant_op.constant([[0, 0], [0, 0], [0, 0], [0, 0]],
+                                      dtype=dtypes.int32)
+      kernel = constant_op.constant(
+          123, shape=[1, 3, 4, 1], dtype=dtypes.float32)
+      self.evaluate(
+          gen_nn_ops.fused_resize_and_pad_conv2d(
+              input=tensor,
+              size=size,
+              paddings=paddings,
+              filter=kernel,
+              mode=mode,
+              strides=strides,
+              padding=padding,
+              resize_align_corners=resize_align_corners))
+
 
 if __name__ == "__main__":
   for index, (input_size_, filter_size_, output_size_, stride_,