Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix failed check in Conv3DBackpropFilterV2.
Passing in a rank-0 `filter_size` causes a check fail and crash,
coming from a `filter_size.vec<>()` call.  Here we check the size
first.

PiperOrigin-RevId: 445517122
  • Loading branch information
cantonios authored and tensorflower-gardener committed Apr 29, 2022
1 parent 03018aa commit 174c509
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/conv_grad_ops_3d.cc
Expand Up @@ -741,6 +741,10 @@ class Conv3DBackpropFilterOp : public OpKernel {
TensorShape filter_shape;
if (takes_shape_) {
const Tensor& filter_sizes = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsVector(filter_sizes.shape()),
errors::InvalidArgument(
"filter_sizes shape must be rank 1 but is rank ",
filter_sizes.shape().dims()));
OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
filter_sizes.vec<int32>(), &filter_shape));
} else {
Expand Down Expand Up @@ -875,6 +879,10 @@ class Conv3DCustomBackpropFilterOp : public OpKernel {
TensorShape filter_shape;
if (takes_shape_) {
const Tensor& filter_sizes = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsVector(filter_sizes.shape()),
errors::InvalidArgument(
"filter_sizes shape must be rank 1 but is rank ",
filter_sizes.shape().dims()));
OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
filter_sizes.vec<int32>(), &filter_shape));
} else {
Expand Down Expand Up @@ -1638,6 +1646,10 @@ class Conv3DBackpropFilterOp<GPUDevice, T> : public OpKernel {
TensorShape filter_shape;
if (takes_shape_) {
const Tensor& filter_sizes = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsVector(filter_sizes.shape()),
errors::InvalidArgument(
"filter_sizes shape must be rank 1 but is rank ",
filter_sizes.shape().dims()));
OP_REQUIRES_OK(context, tensor::MakeShape(filter_sizes, &filter_shape));
} else {
filter_shape = context->input(1).shape();
Expand Down
Expand Up @@ -18,6 +18,7 @@

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker
Expand Down Expand Up @@ -58,6 +59,23 @@ def testGradient(self):
err_tolerance = 1e-3
self.assertLess(err, err_tolerance)

def testBadFilterShape(self):
strides = [1, 1, 1, 1, 1]
padding = "VALID"
tin = constant_op.constant(
.5053710941, shape=[2, 2, 2, 2, 1], dtype=dtypes.float32)
filter_sizes = constant_op.constant(0, shape=[], dtype=dtypes.int32)
out_backprop = constant_op.constant(
.5053710941, shape=[2, 2, 2, 2, 1], dtype=dtypes.float32)

with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
"must be rank 1"):
nn_ops.conv3d_backprop_filter_v2(
input=tin,
filter_sizes=filter_sizes,
out_backprop=out_backprop,
strides=strides,
padding=padding)

if __name__ == "__main__":
test.main()

0 comments on commit 174c509

Please sign in to comment.