Skip to content
Permalink
Browse files Browse the repository at this point in the history
Validate some shape requirements for Conv3DBackpropFilter* and `Con…
…v3DBackpropInput*` ops.

Older versions of Eigen might otherwise crash / produce OOB read on specially crafted inputs.

PiperOrigin-RevId: 369293977
Change-Id: I58f51445a93936d7cf8e616f75de17677df36718
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 19, 2021
1 parent 4814faf commit 8f37b52
Showing 1 changed file with 56 additions and 0 deletions.
56 changes: 56 additions & 0 deletions tensorflow/core/kernels/conv_grad_ops_3d.cc
Expand Up @@ -239,6 +239,20 @@ class Conv3DBackpropInputOp : public OpKernel {
input_shape = context->input(0).shape();
}

OP_REQUIRES(
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
errors::InvalidArgument("input and filter_sizes must have the same "
"number of channels. Got ",
input_shape.dim_size(4), " for input and ",
filter_shape.dim_size(3), " for filter_sizes"));
OP_REQUIRES(
context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
errors::InvalidArgument("out_backprop and filter_sizes must have the "
"same number of channels. Got ",
out_backprop_shape.dim_size(4),
" for out_backprop and ",
filter_shape.dim_size(4), " for filter_sizes"));

ConvBackpropDimensions dims;
OP_REQUIRES_OK(context, ConvBackpropComputeDimensions(
"Conv3DBackpropInputOp", /*num_spatial_dims=*/3,
Expand Down Expand Up @@ -346,6 +360,20 @@ class Conv3DCustomBackpropInputOp : public OpKernel {
input_shape = context->input(0).shape();
}

OP_REQUIRES(
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
errors::InvalidArgument("input and filter_sizes must have the same "
"number of channels. Got ",
input_shape.dim_size(4), " for input and ",
filter_shape.dim_size(3), " for filter_sizes"));
OP_REQUIRES(
context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
errors::InvalidArgument("out_backprop and filter_sizes must have the "
"same number of channels. Got ",
out_backprop_shape.dim_size(4),
" for out_backprop and ",
filter_shape.dim_size(4), " for filter_sizes"));

ConvBackpropDimensions dims;
OP_REQUIRES_OK(context, ConvBackpropComputeDimensions(
"Conv3DBackpropInputOp", /*num_spatial_dims=*/3,
Expand Down Expand Up @@ -696,6 +724,20 @@ class Conv3DBackpropFilterOp : public OpKernel {
filter_shape = context->input(1).shape();
}

OP_REQUIRES(
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
errors::InvalidArgument("input and filter_sizes must have the same "
"number of channels. Got ",
input_shape.dim_size(4), " for input and ",
filter_shape.dim_size(3), " for filter_sizes"));
OP_REQUIRES(
context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
errors::InvalidArgument("out_backprop and filter_sizes must have the "
"same number of channels. Got ",
out_backprop_shape.dim_size(4),
" for out_backprop and ",
filter_shape.dim_size(4), " for filter_sizes"));

ConvBackpropDimensions dims;
OP_REQUIRES_OK(context,
ConvBackpropComputeDimensions(
Expand Down Expand Up @@ -808,6 +850,20 @@ class Conv3DCustomBackpropFilterOp : public OpKernel {
filter_shape = context->input(1).shape();
}

OP_REQUIRES(
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
errors::InvalidArgument("input and filter_sizes must have the same "
"number of channels. Got ",
input_shape.dim_size(4), " for input and ",
filter_shape.dim_size(3), " for filter_sizes"));
OP_REQUIRES(
context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
errors::InvalidArgument("out_backprop and filter_sizes must have the "
"same number of channels. Got ",
out_backprop_shape.dim_size(4),
" for out_backprop and ",
filter_shape.dim_size(4), " for filter_sizes"));

ConvBackpropDimensions dims;
OP_REQUIRES_OK(context,
ConvBackpropComputeDimensions(
Expand Down

0 comments on commit 8f37b52

Please sign in to comment.