【Hackathon 9th No.8】Fix 0-size for as_strided grad and add bound check for as_strided (PaddlePaddle#74860)

* Fix 0-size for as_strided grad and add bound check for as_strided

* fix bugs

* fix bugs

* fix bugs

* refine code

* refine code

* refine code

* refine code to pass coverage ci

Author: ooooo-create

paddle/phi/kernels/stride/as_strided_grad_kernel.cc

@@ -42,6 +42,9 @@ void AsStridedGradKernel(const Context& dev_ctx,
                        phi::StridedTensorFill<data_t>(
                            *input_grad, 0, input_grad);
                      }));
+  if (out_grad.numel() == 0) {
+    return;
+  }
   DenseTensor tmp;
   tmp.set_meta(out_grad.meta());
   AsStridedKernel<Context>(dev_ctx, *input_grad, dims, stride, offset, &tmp);

paddle/phi/kernels/stride/as_strided_kernel.cc

@@ -19,6 +19,26 @@
 COMMON_DECLARE_bool(use_stride_kernel);
 
 namespace phi {
+void ValidateZeroSizeTensorShape(const std::vector<int64_t>& dims,
+                                 const std::vector<int64_t>& strides,
+                                 const DenseTensor& input) {
+  if (input.numel() != 0) {
+    return;
+  }
+  PADDLE_ENFORCE_EQ(dims.size(),
+                    strides.size(),
+                    common::errors::InvalidArgument(
+                        "The size of dims and strides should be equal."));
+  for (size_t i = 0; i < dims.size(); i++) {
+    if (dims[i] == 0) {
+      return;
+    }
+  }
+
+  PADDLE_THROW(common::errors::InvalidArgument(
+      "When input is zero-size tensor, the shape attribute must also be "
+      "zero-size."));
+}
 
 template <typename Context>
 void AsStridedKernel(const Context& dev_ctx,
@@ -36,6 +56,12 @@ void AsStridedKernel(const Context& dev_ctx,
   meta.dims = DDim(dims.data(), static_cast<int>(dims.size()));
   meta.strides = DDim(stride.data(), static_cast<int>(stride.size()));
   meta.offset = offset;
+  ValidateZeroSizeTensorShape(dims, stride, input);
+  PADDLE_ENFORCE_GE(
+      offset,
+      0,
+      common::errors::InvalidArgument(
+          "The offset must be non-negative, but got %d.", offset));
   out->set_meta(meta);
   out->ResetHolder(input.Holder());
   out->ShareInplaceVersionCounterWith(input);

test/legacy_test/test_as_strided.py

@@ -59,5 +59,35 @@ def test_as_strided_backward(self):
                 self.assertEqual((b.grad.numpy() == 1).all().item(), True)
 
 
+class TestAsStrided_ZeroSize(unittest.TestCase):
+    def setUp(self):
+        self.places = get_places()
+
+    def test_as_strided_forward(self):
+        for place in self.places:
+            with base.dygraph.guard(place):
+                a = paddle.to_tensor(
+                    np.random.random([0, 32]).astype('float32')
+                )
+                a.stop_gradient = False
+                b = paddle.as_strided(a, shape=(0, 4), stride=(32, 1))
+                np.testing.assert_equal(b.shape, [0, 4])
+                b.backward(paddle.ones_like(b))
+                np.testing.assert_equal(a.grad.shape, [0, 32])
+
+    def test_as_strided_error(self):
+        for place in self.places:
+            with base.dygraph.guard(place):
+                self.assertRaises(
+                    ValueError,
+                    paddle.as_strided,
+                    x=paddle.to_tensor(
+                        np.random.random([0, 32]).astype('float32')
+                    ),
+                    shape=[3, 4],
+                    stride=[32, 1],
+                )
+
+
 if __name__ == '__main__':
     unittest.main()

test/legacy_test/test_narrow.py

@@ -267,19 +267,18 @@ def setUp(self):
         self.length = 1
 
 
-# TODO(Difers) Address the 0-size issue in the as_strided operator.”
-# class TestPaddleNarrowEmptyTensor(TestNarrowBase):
-#     def setUp(self):
-#         self.input_np = np.empty((0, 4), dtype='float32')
-#         self.input_shape = self.input_np.shape
-#         self.input_dtype = 'float32'
-#         self.op_static = lambda x: paddle.narrow(x, dim=0, start=0, length=0)
-#         self.op_dygraph = lambda x: paddle.narrow(x, dim=0, start=0, length=0)
-#         self.expected = lambda x: x[0:0, :]
-#         self.places = [None, paddle.CPUPlace()]
-#         self.dim = 0
-#         self.start = 0
-#         self.length = 0
+class TestPaddleNarrowEmptyTensor(TestNarrowBase):
+    def setUp(self):
+        self.input_np = np.empty((0, 4), dtype='float32')
+        self.input_shape = self.input_np.shape
+        self.input_dtype = 'float32'
+        self.op_static = lambda x: paddle.narrow(x, dim=0, start=0, length=0)
+        self.op_dygraph = lambda x: paddle.narrow(x, dim=0, start=0, length=0)
+        self.expected = lambda x: x[0:0, :]
+        self.places = [None, paddle.CPUPlace()]
+        self.dim = 0
+        self.start = 0
+        self.length = 0
 
 
 @unittest.skipIf(paddle.device.get_device().startswith("xpu"), "Skip on XPU")