Merge pull request #128296 from vishwakt:fix-conv-zero-size-output-shape

PiperOrigin-RevId: 1005387882

tensorflow/compiler/tests/conv2d_test.py

@@ -404,6 +404,20 @@ class Conv2DBackpropInputTest(xla_test.XLATestCase, parameterized.TestCase):
         data_format_dst=data_format,
         expected=expected_output)
 
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DEmptyInputStride2(self, data_format):
+    # With SAME padding, an input with no rows gives an empty out_backprop.
+    self._VerifyValues(
+        input_sizes=[1, 0, 4, 1],
+        filter_sizes=[3, 3, 1, 1],
+        out_backprop_sizes=[1, 0, 2, 1],
+        strides=[2, 2],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.zeros([1, 0, 4, 1]),
+    )
+
   @parameterized.named_parameters(*DATA_FORMATS)
   def testConv2D1x2FilterStride3Width5(self, data_format):
     expected_output = [1, 2, 0, 2, 4]
@@ -714,6 +728,20 @@ class Conv2DBackpropFilterTest(xla_test.XLATestCase, parameterized.TestCase):
         data_format_dst=data_format,
         expected=expected_output)
 
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DEmptyInput(self, data_format):
+    # With SAME padding, an input with no rows gives an empty out_backprop.
+    self._VerifyValues(
+        input_sizes=[1, 0, 4, 1],
+        filter_sizes=[3, 3, 1, 1],
+        out_backprop_sizes=[1, 0, 4, 1],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.zeros([3, 3, 1, 1]),
+    )
+
   @parameterized.named_parameters(*DATA_FORMATS)
   def testConv2D1x2Filter(self, data_format):
     expected_output = [120, 141]

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -312,6 +312,7 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:tensor_float_32_hdr_lib",
         "@xla//xla:literal_util",
+        "@xla//xla:shape_util",
         "@xla//xla:util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/hlo/builder:xla_builder",

tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc

@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/literal_util.h"
+#include "xla/shape_util.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/bounds_check.h"
@@ -384,6 +385,13 @@ absl::StatusOr<xla::XlaOp> MakeXlaBackpropInputConvOp(
       out_backprop_shape, attrs.dilations, attrs.strides, attrs.padding,
       attrs.data_format, &dims, attrs.explicit_paddings));
 
+  // With an empty `out_backprop` and a stride above 1, the convolution below
+  // doesn't produce `input_shape`. The input gradient is zero then.
+  if (xla::ShapeUtil::IsZeroElementArray(out_backprop_shape)) {
+    return xla::Broadcast(xla::Zero(builder, input_shape.element_type()),
+                          input_shape.dimensions());
+  }
+
   // The input gradients are computed by a convolution of the output
   // gradients and the filter, with some appropriate padding. See the
   // comment at the top of conv_grad_shape_utils.h for details.
@@ -489,6 +497,15 @@ absl::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp(
       grouped_filter_shape, out_backprop_shape, attrs.dilations, attrs.strides,
       attrs.padding, attrs.data_format, &dims, attrs.explicit_paddings));
 
+  // An empty `out_backprop` would be a convolution window with a zero-sized
+  // dimension below, which XLA rejects. The filter gradient is zero then, as
+  // it is for empty activations.
+  if (input_tensor_shape.num_elements() == 0 ||
+      output_tensor_shape.num_elements() == 0) {
+    return xla::Broadcast(xla::Zero(builder, activations_shape.element_type()),
+                          filter_shape.dimensions());
+  }
+
   // Obtain some useful dimensions:
   // The last two dimensions of the filter are the input and output shapes.
   int num_dims = attrs.num_spatial_dims + 2;

tensorflow/core/kernels/conv_grad_filter_ops.cc

@@ -189,8 +189,10 @@ class Conv2DBackpropFilterOp : public OpKernel {
     if (filter_shape.num_elements() == 0) {
       return;
     }
-    // If input is empty, set gradients to zero.
-    if (input.shape().num_elements() == 0) {
+    // If input or out_backprop is empty, set gradients to zero. Otherwise,
+    // cudnn fails with an empty input.
+    if (input.shape().num_elements() == 0 ||
+        out_backprop.shape().num_elements() == 0) {
       functor::SetZeroFunctor<Device, T> f;
       f(context->eigen_device<Device>(), filter_backprop->flat<T>());
       return;
@@ -309,6 +311,14 @@ class Conv2DCustomBackpropFilterOp : public OpKernel {
       return;
     }
 
+    // If shapes are valid but `input` or `out_backprop` is empty,
+    // filter_backprop should be set to all zeros.
+    if (input.NumElements() == 0 || out_backprop.NumElements() == 0) {
+      functor::SetZeroFunctor<Device, T> set_zero;
+      set_zero(context->eigen_device<Device>(), filter_backprop->flat<T>());
+      return;
+    }
+
     int64_t pad_top, pad_bottom;
     int64_t pad_left, pad_right;
     if (padding_ == Padding::EXPLICIT) {

tensorflow/core/kernels/mkl/mkl_conv_ops.cc

@@ -884,7 +884,7 @@ class MklConvOp : public OpKernel {
         MklDnnShape dst_mkl_shape;
         dst_mkl_shape.SetMklTensor(false);
         AllocateOutputSetMklShape(context, kOutputIndex_Dst, &dst_tensor,
-                                  src_tf_shape, dst_mkl_shape, native_format);
+                                  dst_tf_shape, dst_mkl_shape, native_format);
 
         // MklConv2D/3D also outputs converted filter as 2nd output.
         filter_mkl_shape.SetMklTensor(false);

tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py

@@ -866,6 +866,46 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
         op_name=op_name,
     )
 
+  @test_util.run_in_graph_and_eager_modes
+  def testConv2DEmptyBatchOutputShape(self):
+    # With an empty batch, the output still has the shape of the convolution's
+    # output, which differs from the input's here.
+    output = nn_ops.conv2d(
+        array_ops.zeros([0, 4, 4, 3]),
+        array_ops.zeros([3, 3, 3, 5]),
+        strides=[1, 1, 1, 1],
+        padding="VALID",
+    )
+    self.assertEqual(self.evaluate(output).shape, (0, 2, 2, 5))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConv2DBackpropEmptyInput(self):
+    # An empty batch, or an input with no rows under SAME padding, gives an
+    # empty out_backprop, and the gradients are zero.
+    for input_sizes, out_backprop_sizes, stride, padding in (
+        ([0, 4, 4, 3], [0, 2, 2, 2], 1, "VALID"),
+        ([2, 0, 4, 3], [2, 0, 4, 2], 1, "SAME"),
+        ([2, 0, 4, 3], [2, 0, 2, 2], 2, "SAME"),
+    ):
+      strides = [1, stride, stride, 1]
+      out_backprop = array_ops.zeros(out_backprop_sizes)
+      input_grad = gen_nn_ops.conv2d_backprop_input(
+          input_sizes=input_sizes,
+          filter=array_ops.ones([3, 3, 3, 2]),
+          out_backprop=out_backprop,
+          strides=strides,
+          padding=padding,
+      )
+      filter_grad = gen_nn_ops.conv2d_backprop_filter(
+          input=array_ops.ones(input_sizes),
+          filter_sizes=[3, 3, 3, 2],
+          out_backprop=out_backprop,
+          strides=strides,
+          padding=padding,
+      )
+      self.assertAllEqual(np.zeros(input_sizes), self.evaluate(input_grad))
+      self.assertAllEqual(np.zeros([3, 3, 3, 2]), self.evaluate(filter_grad))
+
   @parameterized.named_parameters(*TEST_PARAMS)
   @test_util.run_in_graph_and_eager_modes
   def testConv2D2x2Filter(self, data_format, dtype, use_gpu, op_name):