From 38ba5f2e8705f963717f3aeb99ab7011a3c5b621 Mon Sep 17 00:00:00 2001
From: FlyingQianMM <245467267@qq.com>
Date: Wed, 9 Nov 2022 10:44:10 +0800
Subject: [PATCH] fix ScaleKernel configuration error where input numel is 0
 (#47111)

* fix scale kernel configuration error where input numel is 0

* fix code stype

* add unit test case for scale op when numel of input x is zero

* fix ci codestyle check

* add cpu and gpu unit test case for scale op when numel of input x is zero

* add uninitialized judgment for input of scale
---
 paddle/phi/kernels/cpu/scale_kernel.cc             |  3 +++
 paddle/phi/kernels/gpu/scale_kernel.cu             |  3 +++
 .../paddle/fluid/tests/unittests/test_scale_op.py  | 14 ++++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/paddle/phi/kernels/cpu/scale_kernel.cc b/paddle/phi/kernels/cpu/scale_kernel.cc
index 358d89197ed..53a1391b0cf 100644
--- a/paddle/phi/kernels/cpu/scale_kernel.cc
+++ b/paddle/phi/kernels/cpu/scale_kernel.cc
@@ -40,6 +40,9 @@ void ScaleKernel(const Context& dev_ctx,
   // TODO(chenweihang): now the eigen function here need the dtype of scale,
   // eigen_x, bias should be same, so here need cast for two scalar arg,
   // maybe we declare that the type of scale and bias is T?
+  if (x.numel() <= 0 || (!x.IsInitialized())) {
+    return;
+  }
   paddle::operators::EigenScale<std::decay_t<decltype(dev)>, T>::Eval(
       dev,
       eigen_out,
diff --git a/paddle/phi/kernels/gpu/scale_kernel.cu b/paddle/phi/kernels/gpu/scale_kernel.cu
index 1a574c05494..0efcd0b7063 100644
--- a/paddle/phi/kernels/gpu/scale_kernel.cu
+++ b/paddle/phi/kernels/gpu/scale_kernel.cu
@@ -53,6 +53,9 @@ void ScaleKernel(const Context& dev_ctx,
   inputs.emplace_back(&x);
   outputs.emplace_back(out);
   dev_ctx.template Alloc<T>(out);
+  if (x.numel() <= 0 || (!x.IsInitialized())) {
+    return;
+  }
   phi::funcs::ElementwiseKernel<T>(
       dev_ctx,
       inputs,
diff --git a/python/paddle/fluid/tests/unittests/test_scale_op.py b/python/paddle/fluid/tests/unittests/test_scale_op.py
index 4fec8990537..07be7620a93 100644
--- a/python/paddle/fluid/tests/unittests/test_scale_op.py
+++ b/python/paddle/fluid/tests/unittests/test_scale_op.py
@@ -299,5 +299,19 @@ class TestScaleTripleGradCheck(unittest.TestCase):
             self.func(p)
 
 
+class TestScaleOpZeroNumelVariable(unittest.TestCase):
+    def test_check_zero_numel_cpu(self):
+        paddle.set_device('cpu')
+        data = paddle.ones([0, 1])
+        out = paddle.scale(data, 2)
+        self.assertEqual(out, data)
+
+        if paddle.is_compiled_with_cuda():
+            paddle.set_device('gpu')
+            data = paddle.ones([0, 1])
+            out = paddle.scale(data, 2)
+            self.assertEqual(out, data)
+
+
 if __name__ == "__main__":
     unittest.main()
-- 
GitLab