[Divide by 0 Error] add norm check (#49966)

* [Divide by 0 Error] add norm check * [Divide by 0 Error] fix x AttributeError * [Divide by 0 Error] norm check migrate to c++
PaddlePaddle · Feb 1, 2023 · 5dfddae · 5dfddae
1 parent bdae548
commit 5dfddae
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 0 deletions.
diff --git a/paddle/phi/kernels/cpu/p_norm_kernel.cc b/paddle/phi/kernels/cpu/p_norm_kernel.cc
@@ -61,6 +61,13 @@ void PNormKernel(const Context& dev_ctx,
   int pre, n, post;
   GetDims(xdim, axis, &pre, &n, &post, asvector);
 
+  for (int i = 0; i < xdim.size(); i++) {
+    PADDLE_ENFORCE_LT(0,
+                      xdim[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
+  }
+
   auto* place = dev_ctx.eigen_device();
 
   Eigen::DSizes<int, 3> shape(pre, n, post);

diff --git a/paddle/phi/kernels/gpu/p_norm_kernel.cu b/paddle/phi/kernels/gpu/p_norm_kernel.cu
@@ -105,6 +105,13 @@ void PNormKernel(const Context& dev_ctx,
   std::vector<int> reduce_axis =
       funcs::details::GetReduceDim(axis_dims, xdim.size(), asvector);
 
+  for (int i = 0; i < xdim.size(); i++) {
+    PADDLE_ENFORCE_LT(0,
+                      xdim[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
+  }
+
   using MT = typename dtype::MPTypeTrait<T>::Type;
   if (porder == 0) {
     phi::funcs::ReduceKernel<T, T, kps::AddFunctor, NonzeroFunctor<T>>(

diff --git a/paddle/phi/kernels/xpu/p_norm_kernel.cc b/paddle/phi/kernels/xpu/p_norm_kernel.cc
@@ -55,6 +55,14 @@ void PNormKernel(const Context& dev_ctx,
   int n = 1;
   int t = 1;
   GetDims(xdim, axis, &m, &t, &n, asvector);
+
+  for (int i = 0; i < xdim.size(); i++) {
+    PADDLE_ENFORCE_LT(0,
+                      xdim[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
+  }
+
   x_dim.push_back(m);
   x_dim.push_back(t);
   x_dim.push_back(n);

diff --git a/python/paddle/fluid/tests/unittests/test_norm_all.py b/python/paddle/fluid/tests/unittests/test_norm_all.py
@@ -655,6 +655,15 @@ def err_dtype(p, shape_x, xdtype, out=None):
                 ValueError, paddle.norm, data, p='unspport', axis=[-3, -2, -1]
             )
 
+        with fluid.dygraph.guard():
+            # The size of input in Norm should not be 0.
+            def test_0_size():
+                array = np.array([], dtype=np.float32)
+                x = paddle.to_tensor(np.reshape(array, [0, 0]), dtype='float32')
+                paddle.linalg.norm(x, axis=0)
+
+            self.assertRaises(ValueError, test_0_size)
+
 
 if __name__ == '__main__':
     paddle.enable_static()