Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tensor shape overflow in FusedResizeAndPadConv2D.
Replaced TensorShape constructor by Factory method with status.

PiperOrigin-RevId: 477742686
  • Loading branch information
cantonios authored and tensorflower-gardener committed Sep 29, 2022
1 parent 0baf76b commit d66e1d5
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 4 deletions.
7 changes: 5 additions & 2 deletions tensorflow/core/kernels/conv_ops_fused_image_transform.cc
Expand Up @@ -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) {
Expand Down
4 changes: 2 additions & 2 deletions tensorflow/core/ops/nn_ops.cc
Expand Up @@ -581,7 +581,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")
Expand All @@ -594,7 +594,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);
});

// --------------------------------------------------------------------------
Expand Down
27 changes: 27 additions & 0 deletions tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
Expand Up @@ -3429,6 +3429,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_,
Expand Down

0 comments on commit d66e1d5

Please sign in to comment.