From 311403edbc9816df80274bd1ea8b3c0c0f22c3fa Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Mon, 19 Apr 2021 16:00:40 -0700 Subject: [PATCH] Eliminate a division by 0 in 3D convolutions. Also prevent a CHECK failed introduced in the most recent change. PiperOrigin-RevId: 369322073 Change-Id: I4f609c028f89565fb2b49c3fdd20b63496582bae --- tensorflow/core/kernels/conv_grad_ops_3d.cc | 42 +++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tensorflow/core/kernels/conv_grad_ops_3d.cc b/tensorflow/core/kernels/conv_grad_ops_3d.cc index 8c72d01578d6d..c0b57a7ae5606 100644 --- a/tensorflow/core/kernels/conv_grad_ops_3d.cc +++ b/tensorflow/core/kernels/conv_grad_ops_3d.cc @@ -239,6 +239,14 @@ class Conv3DBackpropInputOp : public OpKernel { input_shape = context->input(0).shape(); } + OP_REQUIRES(context, input_shape.dims() == 5, + errors::InvalidArgument("input tensor must have 5 dimensions")); + OP_REQUIRES( + context, filter_shape.dims() == 5, + errors::InvalidArgument("filter_sizes tensor must have 5 dimensions")); + OP_REQUIRES( + context, out_backprop_shape.dims() == 5, + errors::InvalidArgument("out_backprop tensor must have 5 dimensions")); OP_REQUIRES( context, input_shape.dim_size(4) == filter_shape.dim_size(3), errors::InvalidArgument("input and filter_sizes must have the same " @@ -360,6 +368,14 @@ class Conv3DCustomBackpropInputOp : public OpKernel { input_shape = context->input(0).shape(); } + OP_REQUIRES(context, input_shape.dims() == 5, + errors::InvalidArgument("input tensor must have 5 dimensions")); + OP_REQUIRES( + context, filter_shape.dims() == 5, + errors::InvalidArgument("filter_sizes tensor must have 5 dimensions")); + OP_REQUIRES( + context, out_backprop_shape.dims() == 5, + errors::InvalidArgument("out_backprop tensor must have 5 dimensions")); OP_REQUIRES( context, input_shape.dim_size(4) == filter_shape.dim_size(3), errors::InvalidArgument("input and filter_sizes must have the same " @@ -444,6 +460,11 @@ class Conv3DCustomBackpropInputOp : public OpKernel { // contraction compared to sharding and matmuls. const bool use_parallel_contraction = dims.batch_size == 1; + OP_REQUIRES( + context, work_unit_size > 0, + errors::InvalidArgument("input, filter_sizes and out_backprop tensors " + "must all have at least 1 element")); + const size_t shard_size = use_parallel_contraction ? 1 @@ -724,6 +745,14 @@ class Conv3DBackpropFilterOp : public OpKernel { filter_shape = context->input(1).shape(); } + OP_REQUIRES(context, input_shape.dims() == 5, + errors::InvalidArgument("input tensor must have 5 dimensions")); + OP_REQUIRES( + context, filter_shape.dims() == 5, + errors::InvalidArgument("filter_sizes tensor must have 5 dimensions")); + OP_REQUIRES( + context, out_backprop_shape.dims() == 5, + errors::InvalidArgument("out_backprop tensor must have 5 dimensions")); OP_REQUIRES( context, input_shape.dim_size(4) == filter_shape.dim_size(3), errors::InvalidArgument("input and filter_sizes must have the same " @@ -850,6 +879,14 @@ class Conv3DCustomBackpropFilterOp : public OpKernel { filter_shape = context->input(1).shape(); } + OP_REQUIRES(context, input_shape.dims() == 5, + errors::InvalidArgument("input tensor must have 5 dimensions")); + OP_REQUIRES( + context, filter_shape.dims() == 5, + errors::InvalidArgument("filter_sizes tensor must have 5 dimensions")); + OP_REQUIRES( + context, out_backprop_shape.dims() == 5, + errors::InvalidArgument("out_backprop tensor must have 5 dimensions")); OP_REQUIRES( context, input_shape.dim_size(4) == filter_shape.dim_size(3), errors::InvalidArgument("input and filter_sizes must have the same " @@ -936,6 +973,11 @@ class Conv3DCustomBackpropFilterOp : public OpKernel { const int64 work_unit_size = size_A + size_B + size_C; + OP_REQUIRES( + context, work_unit_size > 0, + errors::InvalidArgument("input, filter_sizes and out_backprop tensors " + "must all have at least 1 element")); + const size_t shard_size = (target_working_set_size + work_unit_size - 1) / work_unit_size;